    {-# 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.IR.SOACS as I hiding (stmPat)
    import Futhark.Internalise.AccurateSizes
    import Futhark.Internalise.Bindings
    import Futhark.Internalise.Entry
    import Futhark.Internalise.Lambdas
    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 -> VisibleTypes -> [E.ValBind] -> m (I.Prog SOACS)
    transformProg :: forall (m :: * -> *).
    MonadFreshNames m =>
    Bool -> VisibleTypes -> [ValBind] -> m (Prog SOACS)
    transformProg Bool
    always_safe VisibleTypes
    types [ValBind]
    vbinds = do
      (OpaqueTypes
    opaques, Stms SOACS
    consts, [FunDef SOACS]
    funs) <-
        Bool
    -> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
    forall (m :: * -> *).
    MonadFreshNames m =>
    Bool
    -> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
    runInternaliseM Bool
    always_safe (VisibleTypes -> [ValBind] -> InternaliseM ()
    internaliseValBinds VisibleTypes
    types [ValBind]
    vbinds)
      Prog SOACS -> m (Prog SOACS)
    forall rep (m :: * -> *).
    (Renameable rep, MonadFreshNames m) =>
    Prog rep -> m (Prog rep)
    I.renameProg (Prog SOACS -> m (Prog SOACS)) -> Prog SOACS -> m (Prog SOACS)
    forall a b. (a -> b) -> a -> b
    $ OpaqueTypes -> Stms SOACS -> [FunDef SOACS] -> Prog SOACS
    forall rep. OpaqueTypes -> Stms rep -> [FunDef rep] -> Prog rep
    I.Prog OpaqueTypes
    opaques Stms SOACS
    consts [FunDef SOACS]
    funs

    internaliseValBinds :: VisibleTypes -> [E.ValBind] -> InternaliseM ()
    internaliseValBinds :: VisibleTypes -> [ValBind] -> InternaliseM ()
    internaliseValBinds VisibleTypes
    types = (ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ((ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ())
    -> (ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ VisibleTypes -> ValBind -> InternaliseM ()
    internaliseValBind VisibleTypes
    types

    internaliseFunName :: VName -> Name
    internaliseFunName :: VName -> Name
    internaliseFunName = [Char] -> Name
    nameFromString ([Char] -> Name) -> (VName -> [Char]) -> VName -> Name
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString

    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

    internaliseValBind :: VisibleTypes -> E.ValBind -> InternaliseM ()
    internaliseValBind :: VisibleTypes -> ValBind -> InternaliseM ()
    internaliseValBind VisibleTypes
    types fb :: ValBind
    fb@(E.ValBind Maybe (Info EntryPoint)
    entry VName
    fname Maybe (TypeExp Exp VName)
    _ (Info ResRetType
    rettype) [TypeParamBase VName]
    tparams [PatBase Info VName ParamType]
    params Exp
    body Maybe DocComment
    _ [AttrInfo VName]
    attrs SrcLoc
    _) = do
      [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 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 :: [Free [] (Param DeclType)]
    all_params = (Param DeclType -> Free [] (Param DeclType))
    -> [Param DeclType] -> [Free [] (Param DeclType)]
    forall a b. (a -> b) -> [a] -> [b]
    map Param DeclType -> Free [] (Param DeclType)
    forall a. a -> Free [] a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Param DeclType]
    [FParam SOACS]
    shapeparams [Free [] (Param DeclType)]
    -> [Free [] (Param DeclType)] -> [Free [] (Param DeclType)]
    forall a. [a] -> [a] -> [a]
    ++ [[Free [] (Param DeclType)]] -> [Free [] (Param DeclType)]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params'
            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 SOACS
    body', [(TypeBase ExtShape Uniqueness, RetAls)]
    rettype') <- InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    -> InternaliseM
         (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m (Result, a) -> m (Body (Rep m), a)
    buildBody (InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
     -> InternaliseM
          (Body (Rep InternaliseM),
           [(TypeBase ExtShape Uniqueness, RetAls)]))
    -> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    -> InternaliseM
         (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
    forall a b. (a -> b) -> a -> b
    $ do
          [SubExp]
    body_res <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp (VName -> [Char]
    baseString VName
    fname [Char] -> [Char] -> [Char]
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "_res") Exp
    body
          ([TypeBase ExtShape Uniqueness]
    rettype', [RetAls]
    retals) <-
            ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness])
    -> ([TypeBase ExtShape Uniqueness], [RetAls])
    -> ([TypeBase ExtShape Uniqueness], [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 Uniqueness] -> [TypeBase ExtShape Uniqueness]
    forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    zeroExts (([TypeBase ExtShape Uniqueness], [RetAls])
     -> ([TypeBase ExtShape Uniqueness], [RetAls]))
    -> ([TypeBase Shape NoUniqueness]
        -> ([TypeBase ExtShape Uniqueness], [RetAls]))
    -> [TypeBase Shape NoUniqueness]
    -> ([TypeBase ExtShape Uniqueness], [RetAls])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [(TypeBase ExtShape Uniqueness, RetAls)]
    -> ([TypeBase ExtShape Uniqueness], [RetAls])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(TypeBase ExtShape Uniqueness, RetAls)]
     -> ([TypeBase ExtShape Uniqueness], [RetAls]))
    -> ([TypeBase Shape NoUniqueness]
        -> [(TypeBase ExtShape Uniqueness, RetAls)])
    -> [TypeBase Shape NoUniqueness]
    -> ([TypeBase ExtShape Uniqueness], [RetAls])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Tree DeclType]
    -> ResRetType
    -> [TypeBase Shape NoUniqueness]
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall shape u.
    [Tree DeclType]
    -> ResRetType
    -> [TypeBase shape u]
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    internaliseReturnType ((Free [] (Param DeclType) -> Tree DeclType)
    -> [Free [] (Param DeclType)] -> [Tree DeclType]
    forall a b. (a -> b) -> [a] -> [b]
    map ((Param DeclType -> DeclType)
    -> Free [] (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) [Free [] (Param DeclType)]
    all_params) ResRetType
    rettype
              ([TypeBase Shape NoUniqueness]
     -> ([TypeBase ExtShape Uniqueness], [RetAls]))
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM ([TypeBase ExtShape Uniqueness], [RetAls])
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    body_res

          Bool -> InternaliseM () -> InternaliseM ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when ([[Free [] (Param DeclType)]] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params') (InternaliseM () -> InternaliseM ())
    -> InternaliseM () -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $
            AppRes -> [SubExp] -> InternaliseM ()
    bindExtSizes (StructType -> [VName] -> AppRes
    E.AppRes (TypeBase Exp Uniqueness -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    E.toStruct (TypeBase Exp Uniqueness -> StructType)
    -> TypeBase Exp Uniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ ResRetType -> TypeBase Exp Uniqueness
    forall dim as. RetTypeBase dim as -> TypeBase dim as
    E.retType ResRetType
    rettype) (ResRetType -> [VName]
    forall dim as. RetTypeBase dim as -> [VName]
    E.retDims ResRetType
    rettype)) [SubExp]
    body_res

          Result
    body_res' <-
            ErrorMsg SubExp -> [ExtType] -> Result -> InternaliseM Result
    ensureResultExtShape ErrorMsg SubExp
    msg ((TypeBase ExtShape Uniqueness -> ExtType)
    -> [TypeBase ExtShape Uniqueness] -> [ExtType]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase ExtShape Uniqueness -> ExtType
    forall shape.
    TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
    I.fromDecl [TypeBase ExtShape Uniqueness]
    rettype') (Result -> InternaliseM Result) -> Result -> InternaliseM Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Result
    subExpsRes [SubExp]
    body_res
          let num_ctx :: Int
    num_ctx = Set Int -> Int
    forall a. Set a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([TypeBase ExtShape Uniqueness] -> Set Int
    forall u. [TypeBase ExtShape u] -> Set Int
    shapeContext [TypeBase ExtShape Uniqueness]
    rettype')
          (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    -> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
            ( Result
    body_res',
              Int
    -> (TypeBase ExtShape Uniqueness, RetAls)
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a. Int -> a -> [a]
    replicate Int
    num_ctx (PrimType -> TypeBase ExtShape Uniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64, RetAls
    forall a. Monoid a => a
    mempty)
                [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a. [a] -> [a] -> [a]
    ++ [TypeBase ExtShape Uniqueness]
    -> [RetAls] -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [TypeBase ExtShape Uniqueness]
    rettype' ((RetAls -> RetAls) -> [RetAls] -> [RetAls]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int -> RetAls -> RetAls
    shiftRetAls Int
    num_ctx) [RetAls]
    retals)
            )

        Attrs
    attrs' <- [AttrInfo VName] -> InternaliseM Attrs
    internaliseAttrs [AttrInfo VName]
    attrs

        let fd :: FunDef SOACS
    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 Uniqueness, RetAls)]
    [(RetType SOACS, RetAls)]
    rettype'
                ((Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList [Free [] (Param DeclType)]
    all_params)
                Body SOACS
    body'

        if [[Free [] (Param DeclType)]] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params'
          then VName -> FunDef SOACS -> InternaliseM ()
    bindConstant VName
    fname FunDef SOACS
    fd
          else
            VName -> FunDef SOACS -> FunInfo -> InternaliseM ()
    bindFunction
              VName
    fname
              FunDef SOACS
    fd
              ( [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
    $ ([Free [] (Param DeclType)] -> [Param DeclType])
    -> [[Free [] (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 ((Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList) [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params',
                (Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList [Free [] (Param DeclType)]
    all_params,
                ([TypeBase ExtShape Uniqueness]
     -> [(TypeBase ExtShape Uniqueness, RetAls)])
    -> Maybe [TypeBase ExtShape Uniqueness]
    -> Maybe [(TypeBase ExtShape Uniqueness, 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 Uniqueness]
    -> [RetAls] -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a b. [a] -> [b] -> [(a, b)]
    `zip` ((TypeBase ExtShape Uniqueness, RetAls) -> RetAls)
    -> [(TypeBase ExtShape Uniqueness, RetAls)] -> [RetAls]
    forall a b. (a -> b) -> [a] -> [b]
    map (TypeBase ExtShape Uniqueness, RetAls) -> RetAls
    forall a b. (a, b) -> b
    snd [(TypeBase ExtShape Uniqueness, RetAls)]
    rettype')
                  (Maybe [TypeBase ExtShape Uniqueness]
     -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)])
    -> ([(SubExp, TypeBase Shape NoUniqueness)]
        -> Maybe [TypeBase ExtShape Uniqueness])
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [TypeBase ExtShape Uniqueness]
    -> [Param DeclType]
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [TypeBase ExtShape Uniqueness]
    forall dec.
    Typed dec =>
    [TypeBase ExtShape Uniqueness]
    -> [Param dec]
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [TypeBase ExtShape Uniqueness]
    forall rt dec.
    (IsRetType rt, Typed dec) =>
    [rt]
    -> [Param dec]
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [rt]
    applyRetType (((TypeBase ExtShape Uniqueness, RetAls)
     -> TypeBase ExtShape Uniqueness)
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map (TypeBase ExtShape Uniqueness, RetAls)
    -> TypeBase ExtShape Uniqueness
    forall a b. (a, b) -> a
    fst [(TypeBase ExtShape Uniqueness, RetAls)]
    rettype') ((Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList [Free [] (Param DeclType)]
    all_params)
              )

      case Maybe (Info EntryPoint)
    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 ()
      where
        zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    zeroExts [TypeBase ExtShape u]
    ts = [TypeBase ExtShape u]
    -> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    forall u.
    [TypeBase ExtShape u]
    -> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    generaliseExtTypes [TypeBase ExtShape u]
    ts [TypeBase ExtShape u]
    ts

    generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM ()
    generateEntryPoint :: VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
    generateEntryPoint VisibleTypes
    types (E.EntryPoint [EntryParam]
    e_params EntryType
    e_rettype) ValBind
    vb = do
      let (E.ValBind Maybe (Info EntryPoint)
    _ VName
    ofname 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 :: [Free [] (Param DeclType)]
    all_params = (Param DeclType -> Free [] (Param DeclType))
    -> [Param DeclType] -> [Free [] (Param DeclType)]
    forall a b. (a -> b) -> [a] -> [b]
    map Param DeclType -> Free [] (Param DeclType)
    forall a. a -> Free [] a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Param DeclType]
    [FParam SOACS]
    shapeparams [Free [] (Param DeclType)]
    -> [Free [] (Param DeclType)] -> [Free [] (Param DeclType)]
    forall a. [a] -> [a] -> [a]
    ++ [[Free [] (Param DeclType)]] -> [Free [] (Param DeclType)]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params'
            ([[TypeBase ExtShape Uniqueness]]
    entry_rettype, [[RetAls]]
    retals) =
              [([TypeBase ExtShape Uniqueness], [RetAls])]
    -> ([[TypeBase ExtShape Uniqueness]], [[RetAls]])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([([TypeBase ExtShape Uniqueness], [RetAls])]
     -> ([[TypeBase ExtShape Uniqueness]], [[RetAls]]))
    -> [([TypeBase ExtShape Uniqueness], [RetAls])]
    -> ([[TypeBase ExtShape Uniqueness]], [[RetAls]])
    forall a b. (a -> b) -> a -> b
    $ ([(TypeBase ExtShape Uniqueness, RetAls)]
     -> ([TypeBase ExtShape Uniqueness], [RetAls]))
    -> [[(TypeBase ExtShape Uniqueness, RetAls)]]
    -> [([TypeBase ExtShape Uniqueness], [RetAls])]
    forall a b. (a -> b) -> [a] -> [b]
    map [(TypeBase ExtShape Uniqueness, RetAls)]
    -> ([TypeBase ExtShape Uniqueness], [RetAls])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([[(TypeBase ExtShape Uniqueness, RetAls)]]
     -> [([TypeBase ExtShape Uniqueness], [RetAls])])
    -> [[(TypeBase ExtShape Uniqueness, RetAls)]]
    -> [([TypeBase ExtShape Uniqueness], [RetAls])]
    forall a b. (a -> b) -> a -> b
    $ [Tree DeclType]
    -> ResRetType -> [[(TypeBase ExtShape Uniqueness, RetAls)]]
    internaliseEntryReturnType ((Free [] (Param DeclType) -> Tree DeclType)
    -> [Free [] (Param DeclType)] -> [Tree DeclType]
    forall a b. (a -> b) -> [a] -> [b]
    map ((Param DeclType -> DeclType)
    -> Free [] (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) [Free [] (Param DeclType)]
    all_params) ResRetType
    rettype
            (EntryPoint
    entry', OpaqueTypes
    opaques) =
              VisibleTypes
    -> Name
    -> [(EntryParam, [Param DeclType])]
    -> (EntryType, [[TypeBase Rank Uniqueness]])
    -> (EntryPoint, OpaqueTypes)
    entryPoint
                VisibleTypes
    types
                (VName -> Name
    baseName VName
    ofname)
                ([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
    $ ([Free [] (Param DeclType)] -> [Param DeclType])
    -> [[Free [] (Param DeclType)]] -> [[Param DeclType]]
    forall a b. (a -> b) -> [a] -> [b]
    map ((Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList) [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params')
                (EntryType
    e_rettype, ([TypeBase ExtShape Uniqueness] -> [TypeBase Rank Uniqueness])
    -> [[TypeBase ExtShape Uniqueness]] -> [[TypeBase Rank Uniqueness]]
    forall a b. (a -> b) -> [a] -> [b]
    map ((TypeBase ExtShape Uniqueness -> TypeBase Rank Uniqueness)
    -> [TypeBase ExtShape Uniqueness] -> [TypeBase Rank Uniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase ExtShape Uniqueness -> TypeBase Rank Uniqueness
    forall shape u.
    ArrayShape shape =>
    TypeBase shape u -> TypeBase Rank u
    I.rankShaped) [[TypeBase ExtShape Uniqueness]]
    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
    $ ([Free [] (Param DeclType)] -> [Param DeclType])
    -> [[Free [] (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 ((Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList) [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params'

        OpaqueTypes -> InternaliseM ()
    addOpaques OpaqueTypes
    opaques

        (Body SOACS
    entry_body, [(TypeBase ExtShape Uniqueness, RetAls)]
    ctx_ts) <- InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    -> InternaliseM
         (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m (Result, a) -> m (Body (Rep m), a)
    buildBody (InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
     -> InternaliseM
          (Body (Rep InternaliseM),
           [(TypeBase ExtShape Uniqueness, RetAls)]))
    -> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    -> InternaliseM
         (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
    forall a b. (a -> b) -> a -> b
    $ do
          -- Special case the (rare) situation where the entry point is
          -- not a function.
          Maybe [SubExp]
    maybe_const <- VName -> InternaliseM (Maybe [SubExp])
    lookupConst VName
    ofname
          [SubExp]
    vals <- case Maybe [SubExp]
    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 ->
              [Char] -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
    funcall [Char]
    "entry_result" (VName -> QualName VName
    forall v. v -> QualName v
    E.qualName VName
    ofname) [SubExp]
    args
          [SubExp]
    ctx <-
            [TypeBase ExtShape Uniqueness] -> [[SubExp]] -> [SubExp]
    forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
    extractShapeContext ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
    forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    zeroExts ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness])
    -> [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> a -> b
    $ [[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[TypeBase ExtShape Uniqueness]]
    entry_rettype)
              ([[SubExp]] -> [SubExp])
    -> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (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 ((TypeBase Shape NoUniqueness -> [SubExp])
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM [SubExp]
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (InternaliseM (TypeBase Shape NoUniqueness)
     -> InternaliseM [SubExp])
    -> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> SubExp
    -> InternaliseM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType) [SubExp]
    vals
          (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    -> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([SubExp] -> Result
    subExpsRes ([SubExp] -> Result) -> [SubExp] -> Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    ctx [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [SubExp]
    vals, (SubExp -> (TypeBase ExtShape Uniqueness, RetAls))
    -> [SubExp] -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((TypeBase ExtShape Uniqueness, RetAls)
    -> SubExp -> (TypeBase ExtShape Uniqueness, RetAls)
    forall a b. a -> b -> a
    const (PrimType -> TypeBase ExtShape Uniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64, RetAls
    forall a. Monoid a => a
    mempty)) [SubExp]
    ctx)

        Attrs
    attrs' <- [AttrInfo VName] -> InternaliseM Attrs
    internaliseAttrs [AttrInfo VName]
    attrs
        let num_ctx :: Int
    num_ctx = [(TypeBase ExtShape Uniqueness, RetAls)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(TypeBase ExtShape Uniqueness, RetAls)]
    ctx_ts
        FunDef SOACS -> InternaliseM ()
    addFunDef (FunDef SOACS -> InternaliseM ())
    -> FunDef SOACS -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $
          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
            (EntryPoint -> Maybe EntryPoint
    forall a. a -> Maybe a
    Just EntryPoint
    entry')
            Attrs
    attrs'
            (Name
    "entry_" Name -> Name -> Name
    forall a. Semigroup a => a -> a -> a
    <> VName -> Name
    baseName VName
    ofname)
            ( [(TypeBase ExtShape Uniqueness, RetAls)]
    ctx_ts
                [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a. [a] -> [a] -> [a]
    ++ [TypeBase ExtShape Uniqueness]
    -> [RetAls] -> [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip
                  ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
    forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    zeroExts ([[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[TypeBase ExtShape Uniqueness]]
    entry_rettype))
                  ((RetAls -> RetAls) -> [RetAls] -> [RetAls]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int -> RetAls -> RetAls
    shiftRetAls Int
    num_ctx) ([RetAls] -> [RetAls]) -> [RetAls] -> [RetAls]
    forall a b. (a -> b) -> a -> b
    $ [[RetAls]] -> [RetAls]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[RetAls]]
    retals)
            )
            ([Param DeclType]
    [FParam SOACS]
    shapeparams [Param DeclType] -> [Param DeclType] -> [Param DeclType]
    forall a. [a] -> [a] -> [a]
    ++ ([Free [] (Param DeclType)] -> [Param DeclType])
    -> [[Free [] (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 ((Free [] (Param DeclType) -> [Param DeclType])
    -> [Free [] (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 Free [] (Param DeclType) -> [Param DeclType]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList) [[Free [] (Param DeclType)]]
    [[Tree (FParam SOACS)]]
    params')
            Body SOACS
    entry_body
      where
        zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    zeroExts [TypeBase ExtShape u]
    ts = [TypeBase ExtShape u]
    -> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    forall u.
    [TypeBase ExtShape u]
    -> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
    generaliseExtTypes [TypeBase ExtShape u]
    ts [TypeBase ExtShape u]
    ts

    internaliseBody :: String -> E.Exp -> InternaliseM (Body SOACS)
    internaliseBody :: [Char] -> Exp -> InternaliseM (Body SOACS)
    internaliseBody [Char]
    desc Exp
    e =
      InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (InternaliseM Result -> InternaliseM (Body (Rep InternaliseM)))
    -> InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Result
    subExpsRes ([SubExp] -> Result)
    -> InternaliseM [SubExp] -> InternaliseM Result
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "_res") Exp
    e

    bodyFromStms ::
      InternaliseM (Result, a) ->
      InternaliseM (Body SOACS, a)
    bodyFromStms :: forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
    bodyFromStms InternaliseM (Result, a)
    m = do
      ((Result
    res, a
    a), Stms SOACS
    stms) <- InternaliseM (Result, a)
    -> InternaliseM ((Result, 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 (Result, a)
    m
      (,a
    a) (Body SOACS -> (Body SOACS, a))
    -> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS, a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Stms (Rep InternaliseM)
    -> Result -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    Stms (Rep m) -> Result -> m (Body (Rep m))
    mkBodyM Stms (Rep InternaliseM)
    Stms SOACS
    stms Result
    res

    -- | Only returns those pattern names that are not used in the pattern
    -- itself (the "non-existential" part, you could say).
    letValExp :: String -> I.Exp SOACS -> InternaliseM [VName]
    letValExp :: [Char] -> Exp SOACS -> InternaliseM [VName]
    letValExp [Char]
    name Exp SOACS
    e = do
      [ExtType]
    e_t <- Exp SOACS -> InternaliseM [ExtType]
    forall rep (m :: * -> *).
    (HasScope rep m, TypedOp (OpC rep)) =>
    Exp rep -> m [ExtType]
    expExtType Exp SOACS
    e
      [VName]
    names <- Int -> InternaliseM VName -> InternaliseM [VName]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM ([ExtType] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [ExtType]
    e_t) (InternaliseM VName -> InternaliseM [VName])
    -> InternaliseM VName -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ [Char] -> InternaliseM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    name
      [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName]
    names Exp (Rep InternaliseM)
    Exp SOACS
    e
      let ctx :: Set Int
    ctx = [ExtType] -> Set Int
    forall u. [TypeBase ExtShape u] -> Set Int
    shapeContext [ExtType]
    e_t
      [VName] -> InternaliseM [VName]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([VName] -> InternaliseM [VName])
    -> [VName] -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ ((VName, Int) -> VName) -> [(VName, Int)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, Int) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, Int)] -> [VName]) -> [(VName, Int)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ ((VName, Int) -> Bool) -> [(VName, Int)] -> [(VName, Int)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((Int -> Set Int -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.notMember` Set Int
    ctx) (Int -> Bool) -> ((VName, Int) -> Int) -> (VName, Int) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Int) -> Int
    forall a b. (a, b) -> b
    snd) ([(VName, Int)] -> [(VName, Int)])
    -> [(VName, Int)] -> [(VName, Int)]
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [Int] -> [(VName, Int)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    names [Int
    0 ..]

    letValExp' :: String -> I.Exp SOACS -> InternaliseM [SubExp]
    letValExp' :: [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    _ (BasicOp (SubExp SubExp
    se)) = [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp
    se]
    letValExp' [Char]
    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
    <$> [Char] -> Exp SOACS -> InternaliseM [VName]
    letValExp [Char]
    name Exp SOACS
    ses

    internaliseAppExp :: String -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp]
    internaliseAppExp :: [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
    internaliseAppExp [Char]
    desc AppRes
    _ (E.Index Exp
    e SliceBase Info VName
    idxs SrcLoc
    _) = do
      [VName]
    vs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "indexed" Exp
    e
      [SubExp]
    dims <- case [VName]
    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 NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    v
      ([DimIndex SubExp]
    idxs', Certs
    cs) <- [SubExp]
    -> SliceBase Info VName -> InternaliseM ([DimIndex SubExp], Certs)
    internaliseSlice [SubExp]
    dims SliceBase Info VName
    idxs
      let index :: VName -> InternaliseM (Exp SOACS)
    index VName
    v = do
            TypeBase Shape NoUniqueness
    v_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    v
            Exp SOACS -> InternaliseM (Exp SOACS)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp SOACS -> InternaliseM (Exp SOACS))
    -> Exp SOACS -> InternaliseM (Exp SOACS)
    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
    $ VName -> Slice SubExp -> BasicOp
    I.Index VName
    v (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
    fullSlice TypeBase Shape NoUniqueness
    v_t [DimIndex SubExp]
    idxs'
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a. Certs -> InternaliseM a -> InternaliseM a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying Certs
    cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ (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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    desc (Exp SOACS -> InternaliseM SubExp)
    -> (VName -> InternaliseM (Exp SOACS))
    -> VName
    -> InternaliseM SubExp
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< VName -> InternaliseM (Exp SOACS)
    index) [VName]
    vs
    internaliseAppExp [Char]
    desc AppRes
    _ (E.Range Exp
    start Maybe Exp
    maybe_second Inclusiveness Exp
    end SrcLoc
    _) = do
      SubExp
    start' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "range_start" Exp
    start
      SubExp
    end' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "range_end" (Exp -> InternaliseM SubExp) -> Exp -> InternaliseM SubExp
    forall a b. (a -> b) -> a -> b
    $ case Inclusiveness Exp
    end of
        DownToExclusive Exp
    e -> Exp
    e
        ToInclusive Exp
    e -> Exp
    e
        UpToExclusive Exp
    e -> Exp
    e
      Maybe SubExp
    maybe_second' <-
        (Exp -> InternaliseM SubExp)
    -> Maybe Exp -> InternaliseM (Maybe SubExp)
    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) -> Maybe a -> f (Maybe b)
    traverse ([Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "range_second") Maybe Exp
    maybe_second

      -- Construct an error message in case the range is invalid.
      let conv :: SubExp -> InternaliseM SubExp
    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
      SubExp
    start'_i64 <- SubExp -> InternaliseM SubExp
    conv SubExp
    start'
      SubExp
    end'_i64 <- SubExp -> InternaliseM SubExp
    conv SubExp
    end'
      Maybe SubExp
    maybe_second'_i64 <- (SubExp -> InternaliseM SubExp)
    -> Maybe SubExp -> InternaliseM (Maybe SubExp)
    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) -> Maybe a -> f (Maybe b)
    traverse SubExp -> InternaliseM SubExp
    conv Maybe SubExp
    maybe_second'
      let errmsg :: ErrorMsg SubExp
    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."]

      (IntType
    it, CmpOp
    lt_op) <-
        case Exp -> StructType
    E.typeOf Exp
    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 :: SubExp
    one = IntType -> Integer -> SubExp
    intConst IntType
    it Integer
    1
          negone :: SubExp
    negone = IntType -> Integer -> SubExp
    intConst IntType
    it (-Integer
    1)
          default_step :: SubExp
    default_step = case Inclusiveness Exp
    end of
            DownToExclusive {} -> SubExp
    negone
            ToInclusive {} -> SubExp
    one
            UpToExclusive {} -> SubExp
    one

      (SubExp
    step, SubExp
    step_zero) <- case Maybe SubExp
    maybe_second' of
        Just SubExp
    second' -> do
          SubExp
    subtracted_step <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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'
          SubExp
    step_zero <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "step_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
    I.CmpOp (PrimType -> CmpOp
    I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
    forall a b. (a -> b) -> a -> b
    $ IntType -> PrimType
    IntType IntType
    it) SubExp
    start' SubExp
    second'
          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    subtracted_step, SubExp
    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)

      SubExp
    step_sign <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "s_sign" (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
    $ UnOp -> SubExp -> BasicOp
    I.UnOp (IntType -> UnOp
    I.SSignum IntType
    it) SubExp
    step
      SubExp
    step_sign_i64 <- IntType -> SubExp -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntS IntType
    Int64 SubExp
    step_sign

      SubExp
    bounds_invalid_downwards <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "bounds_invalid_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 CmpOp
    lt_op SubExp
    start' SubExp
    end'
      SubExp
    bounds_invalid_upwards <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "bounds_invalid_upwards" (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
    lt_op SubExp
    end' SubExp
    start'

      (SubExp
    distance, SubExp
    step_wrong_dir, SubExp
    bounds_invalid) <- case Inclusiveness Exp
    end of
        DownToExclusive {} -> do
          SubExp
    step_wrong_dir <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
          SubExp
    distance <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "distance" (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
    Sub IntType
    it Overflow
    I.OverflowWrap) SubExp
    start' SubExp
    end'
          SubExp
    distance_i64 <- IntType -> SubExp -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntS IntType
    Int64 SubExp
    distance
          (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    distance_i64, SubExp
    step_wrong_dir, SubExp
    bounds_invalid_downwards)
        UpToExclusive {} -> do
          SubExp
    step_wrong_dir <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
          SubExp
    distance <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "distance" (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
    Sub IntType
    it Overflow
    I.OverflowWrap) SubExp
    end' SubExp
    start'
          SubExp
    distance_i64 <- IntType -> SubExp -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntS IntType
    Int64 SubExp
    distance
          (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    distance_i64, SubExp
    step_wrong_dir, SubExp
    bounds_invalid_upwards)
        ToInclusive {} -> do
          SubExp
    downwards <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
          SubExp
    distance_downwards_exclusive <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "distance_downwards_exclusive" (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
    Sub IntType
    it Overflow
    I.OverflowWrap) SubExp
    start' SubExp
    end'
          SubExp
    distance_upwards_exclusive <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "distance_upwards_exclusive" (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
    Sub IntType
    it Overflow
    I.OverflowWrap) SubExp
    end' SubExp
    start'

          SubExp
    bounds_invalid <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "bounds_invalid"
              (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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    downwards)
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    bounds_invalid_downwards])
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    bounds_invalid_upwards])
          SubExp
    distance_exclusive <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "distance_exclusive"
              (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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    downwards)
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    distance_downwards_exclusive])
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    distance_upwards_exclusive])
          SubExp
    distance_exclusive_i64 <- IntType -> SubExp -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntS IntType
    Int64 SubExp
    distance_exclusive
          SubExp
    distance <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "distance" (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
    Add IntType
    Int64 Overflow
    I.OverflowWrap)
                  SubExp
    distance_exclusive_i64
                  (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
          (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    distance, Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    False, SubExp
    bounds_invalid)

      SubExp
    step_invalid <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "step_invalid" (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.LogOr SubExp
    step_wrong_dir SubExp
    step_zero

      SubExp
    invalid <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "range_invalid" (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.LogOr SubExp
    step_invalid SubExp
    bounds_invalid
      SubExp
    valid <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "valid" (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
    invalid
      Certs
    cs <- [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "range_valid_c" SubExp
    valid ErrorMsg SubExp
    errmsg

      SubExp
    step_i64 <- IntType -> SubExp -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntS IntType
    Int64 SubExp
    step
      SubExp
    pos_step <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "pos_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
    Mul IntType
    Int64 Overflow
    I.OverflowWrap) SubExp
    step_i64 SubExp
    step_sign_i64

      SubExp
    num_elems <-
        Certs -> InternaliseM SubExp -> InternaliseM SubExp
    forall a. Certs -> InternaliseM a -> InternaliseM a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying Certs
    cs (InternaliseM SubExp -> InternaliseM SubExp)
    -> InternaliseM SubExp -> InternaliseM SubExp
    forall a b. (a -> b) -> a -> b
    $
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "num_elems" (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 -> Safety -> BinOp
    SDivUp IntType
    Int64 Safety
    I.Unsafe) SubExp
    distance SubExp
    pos_step

      SubExp
    se <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    desc (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 -> SubExp -> SubExp -> IntType -> BasicOp
    I.Iota SubExp
    num_elems SubExp
    start' SubExp
    step IntType
    it)
      [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp
    se]
    internaliseAppExp [Char]
    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
          [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    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 = [Char] -> Name
    nameFromString ([Char] -> Name) -> [Char] -> Name
    forall a b. (a -> b) -> a -> b
    $ Name -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString (Name -> [Char]) -> Name -> [Char]
    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
    qfname
              arg_desc :: [Char]
    arg_desc = Name -> [Char]
    nameToString Name
    fname [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_arg"

          -- Some functions are magical (overloaded) and we handle that here.
          case () of
            ()
              -- Short-circuiting operators are magical.
              | VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qfname) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag,
                VName -> [Char]
    baseString (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qfname) [Char] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Char]
    "&&",
                [(Exp
    x, Maybe VName
    _), (Exp
    y, Maybe VName
    _)] <- [(Exp, Maybe VName)]
    args ->
                  [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    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 NoUniqueness -> StructType
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
    -> ScalarTypeBase Exp NoUniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ PrimType -> ScalarTypeBase Exp NoUniqueness
    forall dim u. PrimType -> ScalarTypeBase dim u
    E.Prim PrimType
    E.Bool) [])
              | VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qfname) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag,
                VName -> [Char]
    baseString (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qfname) [Char] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Char]
    "||",
                [(Exp
    x, Maybe VName
    _), (Exp
    y, Maybe VName
    _)] <- [(Exp, Maybe VName)]
    args ->
                  [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    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 NoUniqueness -> StructType
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
    -> ScalarTypeBase Exp NoUniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ PrimType -> ScalarTypeBase Exp NoUniqueness
    forall dim u. PrimType -> ScalarTypeBase dim u
    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
    -> [Char]
    -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
    isOverloadedFunction QualName VName
    qfname [Char]
    desc -> do
                  let prepareArg :: (Exp, b) -> InternaliseM (StructType, [SubExp])
    prepareArg (Exp
    arg, b
    _) =
                        (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    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
    <$> [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "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 [Char] -> InternaliseM [SubExp]
    internalise <- QualName VName -> [Exp] -> Maybe ([Char] -> 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) ->
                  [Char] -> InternaliseM [SubExp]
    internalise [Char]
    desc
              | VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qfname) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag,
                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]
                  [[SubExp]]
    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 ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
    internaliseArg [Char]
    arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
    forall a. [a] -> [a]
    reverse [(Exp, Maybe VName)]
    args)
                  let args'' :: [(SubExp, Diet)]
    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'
                  [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    desc (Exp SOACS -> InternaliseM [SubExp])
    -> Exp SOACS -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Name
    -> [(SubExp, Diet)]
    -> [(RetType SOACS, RetAls)]
    -> Safety
    -> Exp SOACS
    forall rep.
    Name
    -> [(SubExp, Diet)] -> [(RetType rep, RetAls)] -> Safety -> Exp rep
    I.Apply Name
    fname [(SubExp, Diet)]
    args'' [(PrimType -> TypeBase ExtShape Uniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    rettype, RetAls
    forall a. Monoid a => a
    mempty)] Safety
    Safe
              | Bool
    otherwise -> do
                  [SubExp]
    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 ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
    internaliseArg [Char]
    arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
    forall a. [a] -> [a]
    reverse [(Exp, Maybe VName)]
    args)
                  [Char] -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
    funcall [Char]
    desc QualName VName
    qfname [SubExp]
    args'
    internaliseAppExp [Char]
    desc AppRes
    _ (E.LetPat [SizeBinder VName]
    sizes PatBase Info VName StructType
    pat Exp
    e Exp
    body SrcLoc
    _) =
      [Char]
    -> [SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> InternaliseM [SubExp]
    -> InternaliseM [SubExp]
    forall a.
    [Char]
    -> [SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> InternaliseM a
    -> InternaliseM a
    internalisePat [Char]
    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
    $ [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    body
    internaliseAppExp [Char]
    _ AppRes
    _ (E.LetFun VName
    ofname ([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 [Char]
    desc AppRes
    _ (E.Loop [VName]
    sparams PatBase Info VName ParamType
    mergepat LoopInitBase Info VName
    loopinit LoopFormBase Info VName
    form Exp
    loopbody SrcLoc
    _) = do
      [SubExp]
    ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "loop_init" (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ LoopInitBase Info VName -> Exp
    loopInitExp LoopInitBase Info VName
    loopinit
      ((Body SOACS
    loopbody', (LoopForm
    form', [Param DeclType]
    shapepat, [Param DeclType]
    mergepat', [SubExp]
    mergeinit')), Stms SOACS
    initstms) <-
        InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         ((Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])),
          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
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
     -> InternaliseM
          ((Body SOACS,
            (LoopForm, [Param DeclType], [Param DeclType], [SubExp])),
           Stms (Rep InternaliseM)))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         ((Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])),
          Stms (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    -> LoopFormBase Info VName
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    handleForm [SubExp]
    ses LoopFormBase Info VName
    form

      Stms (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep InternaliseM)
    Stms SOACS
    initstms
      [TypeBase Shape NoUniqueness]
    mergeinit_ts' <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    mergeinit'

      [SubExp]
    ctxinit <- [VName]
    -> [FParam SOACS]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [SubExp]
    argShapes ((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]
    shapepat) [Param DeclType]
    [FParam SOACS]
    mergepat' [TypeBase Shape NoUniqueness]
    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]
    args = [SubExp]
    ctxinit [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [SubExp]
    mergeinit'
      [SubExp]
    args' <-
        ErrorMsg SubExp
    -> [VName]
    -> [TypeBase Shape NoUniqueness]
    -> [SubExp]
    -> InternaliseM [SubExp]
    forall u.
    Typed (TypeBase Shape u) =>
    ErrorMsg SubExp
    -> [VName]
    -> [TypeBase Shape u]
    -> [SubExp]
    -> InternaliseM [SubExp]
    ensureArgShapes
          ErrorMsg SubExp
    "initial loop values have right shape"
          ((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]
    shapepat)
          ((Param DeclType -> TypeBase Shape NoUniqueness)
    -> [Param DeclType] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    paramType ([Param DeclType] -> [TypeBase Shape NoUniqueness])
    -> [Param DeclType] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ [Param DeclType]
    shapepat [Param DeclType] -> [Param DeclType] -> [Param DeclType]
    forall a. [a] -> [a] -> [a]
    ++ [Param DeclType]
    mergepat')
          [SubExp]
    args

      let dropCond :: [VName] -> [VName]
    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)]
    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 :: [TypeBase Shape NoUniqueness]
    merge_ts = ((Param DeclType, SubExp) -> TypeBase Shape NoUniqueness)
    -> [(Param DeclType, SubExp)] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    I.paramType (Param DeclType -> TypeBase Shape NoUniqueness)
    -> ((Param DeclType, SubExp) -> Param DeclType)
    -> (Param DeclType, SubExp)
    -> TypeBase Shape NoUniqueness
    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
      Body SOACS
    loopbody'' <-
        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 DeclType] -> Scope SOACS
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams (((Param DeclType, SubExp) -> Param DeclType)
    -> [(Param DeclType, SubExp)] -> [Param DeclType]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param DeclType, SubExp) -> Param DeclType
    forall a b. (a, b) -> a
    fst [(Param DeclType, SubExp)]
    merge) Scope SOACS -> Scope SOACS -> Scope SOACS
    forall a. Semigroup a => a -> a -> a
    <> LoopForm -> Scope SOACS
    forall rep. LoopForm -> Scope rep
    scopeOfLoopForm LoopForm
    form') (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
    -> (InternaliseM Result -> InternaliseM (Body SOACS))
    -> InternaliseM Result
    -> InternaliseM (Body SOACS)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
    InternaliseM Result -> InternaliseM (Body SOACS)
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (InternaliseM Result -> InternaliseM (Body SOACS))
    -> InternaliseM Result -> InternaliseM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $
          ([SubExp] -> Result)
    -> InternaliseM [SubExp] -> InternaliseM Result
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [SubExp] -> Result
    subExpsRes
            (InternaliseM [SubExp] -> InternaliseM Result)
    -> (Result -> InternaliseM [SubExp])
    -> Result
    -> InternaliseM Result
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ErrorMsg SubExp
    -> [VName]
    -> [TypeBase Shape NoUniqueness]
    -> [SubExp]
    -> InternaliseM [SubExp]
    forall u.
    Typed (TypeBase Shape u) =>
    ErrorMsg SubExp
    -> [VName]
    -> [TypeBase Shape u]
    -> [SubExp]
    -> InternaliseM [SubExp]
    ensureArgShapes
              ErrorMsg SubExp
    "shape of loop result does not match shapes in loop parameter"
              (((Param DeclType, SubExp) -> VName)
    -> [(Param DeclType, SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName (Param DeclType -> VName)
    -> ((Param DeclType, SubExp) -> Param DeclType)
    -> (Param DeclType, SubExp)
    -> VName
    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)
              [TypeBase Shape NoUniqueness]
    merge_ts
            ([SubExp] -> InternaliseM [SubExp])
    -> (Result -> [SubExp]) -> Result -> InternaliseM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (SubExpRes -> SubExp) -> Result -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp
            (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
    loopbody'

      Attrs
    attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks InternaliseEnv -> Attrs
    envAttrs
      (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var ([VName] -> [SubExp])
    -> ([VName] -> [VName]) -> [VName] -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [VName] -> [VName]
    dropCond
        ([VName] -> [SubExp])
    -> InternaliseM [VName] -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Attrs -> InternaliseM [VName] -> InternaliseM [VName]
    forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
    attributing
          Attrs
    attrs
          ([Char] -> Exp SOACS -> InternaliseM [VName]
    letValExp [Char]
    desc ([(FParam SOACS, SubExp)] -> LoopForm -> Body SOACS -> Exp SOACS
    forall rep.
    [(FParam rep, SubExp)] -> LoopForm -> Body rep -> Exp rep
    I.Loop [(Param DeclType, SubExp)]
    [(FParam SOACS, SubExp)]
    merge LoopForm
    form' Body SOACS
    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
    $ [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "loopres" Exp
    loopbody

        forLoop :: [Param DeclType]
    -> [Param DeclType]
    -> [SubExp]
    -> VName
    -> [(Param (TypeBase Shape NoUniqueness), 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 NoUniqueness), VName)]
    loopvars LoopForm
    form' =
          InternaliseM
      (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
    bodyFromStms (InternaliseM
       (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
     -> InternaliseM
          (Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> (InternaliseM
          (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
        -> InternaliseM
             (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> InternaliseM
         (Result, (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
         (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Result, (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
       (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
     -> InternaliseM
          (Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> InternaliseM
         (Result, (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 NoUniqueness), VName)]
    -> ((Param (TypeBase Shape NoUniqueness), VName)
        -> InternaliseM ())
    -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [(Param (TypeBase Shape NoUniqueness), VName)]
    loopvars (((Param (TypeBase Shape NoUniqueness), VName) -> InternaliseM ())
     -> InternaliseM ())
    -> ((Param (TypeBase Shape NoUniqueness), VName)
        -> InternaliseM ())
    -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ \(Param (TypeBase Shape NoUniqueness)
    p, VName
    arr) ->
              [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName Param (TypeBase Shape NoUniqueness)
    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)]
            [SubExp]
    ses <- InternaliseM [SubExp]
    loopBody
            [TypeBase Shape NoUniqueness]
    sets <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    ses
            [SubExp]
    shapeargs <- [VName]
    -> [FParam SOACS]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [SubExp]
    argShapes ((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]
    shapepat) [Param DeclType]
    [FParam SOACS]
    mergepat' [TypeBase Shape NoUniqueness]
    sets
            (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
              ( [SubExp] -> Result
    subExpsRes ([SubExp] -> Result) -> [SubExp] -> Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    shapeargs [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [SubExp]
    ses,
                ( LoopForm
    form',
                  [Param DeclType]
    shapepat,
                  [Param DeclType]
    mergepat',
                  [SubExp]
    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
          [VName]
    arr' <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "for_in_arr" Exp
    arr
          [TypeBase Shape NoUniqueness]
    arr_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arr'
          let w :: SubExp
    w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    arr_ts

          VName
    i <- [Char] -> InternaliseM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "i"

          [TypeBase Shape NoUniqueness]
    ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    mergeinit
          [TypeParamBase VName]
    -> PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([FParam SOACS]
        -> [FParam SOACS]
        -> InternaliseM
             (Body SOACS,
              (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a.
    [TypeParamBase VName]
    -> PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
    -> InternaliseM a
    bindingLoopParams [TypeParamBase VName]
    sparams' PatBase Info VName ParamType
    mergepat [TypeBase Shape NoUniqueness]
    ts (([FParam SOACS]
      -> [FParam SOACS]
      -> InternaliseM
           (Body SOACS,
            (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
     -> InternaliseM
          (Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> ([FParam SOACS]
        -> [FParam 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
    $ \[FParam SOACS]
    shapepat [FParam SOACS]
    mergepat' ->
            [PatBase Info VName ParamType]
    -> [TypeBase Shape NoUniqueness]
    -> ([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 NoUniqueness]
    -> ([LParam SOACS] -> InternaliseM a)
    -> InternaliseM a
    bindingLambdaParams [Diet -> StructType -> ParamType
    forall u. Diet -> TypeBase Exp u -> 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 NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType [TypeBase Shape NoUniqueness]
    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 NoUniqueness), VName)]
    loopvars = [Param (TypeBase Shape NoUniqueness)]
    -> [VName] -> [(Param (TypeBase Shape NoUniqueness), VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param (TypeBase Shape NoUniqueness)]
    [LParam SOACS]
    x_params [VName]
    arr'
              [Param DeclType]
    -> [Param DeclType]
    -> [SubExp]
    -> VName
    -> [(Param (TypeBase Shape NoUniqueness), 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 NoUniqueness), 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
          SubExp
    num_iterations' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "upper_bound" Exp
    num_iterations
          TypeBase Shape NoUniqueness
    num_iterations_t <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType SubExp
    num_iterations'
          IntType
    it <- case TypeBase Shape NoUniqueness
    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 NoUniqueness
    _ -> [Char] -> InternaliseM IntType
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "internaliseExp Loop: invalid type"

          [TypeBase Shape NoUniqueness]
    ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    mergeinit
          [TypeParamBase VName]
    -> PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([FParam SOACS]
        -> [FParam SOACS]
        -> InternaliseM
             (Body SOACS,
              (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a.
    [TypeParamBase VName]
    -> PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
    -> InternaliseM a
    bindingLoopParams [TypeParamBase VName]
    sparams' PatBase Info VName ParamType
    mergepat [TypeBase Shape NoUniqueness]
    ts (([FParam SOACS]
      -> [FParam SOACS]
      -> InternaliseM
           (Body SOACS,
            (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
     -> InternaliseM
          (Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> ([FParam SOACS]
        -> [FParam 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
    $ \[FParam SOACS]
    shapepat [FParam SOACS]
    mergepat' ->
            [Param DeclType]
    -> [Param DeclType]
    -> [SubExp]
    -> VName
    -> [(Param (TypeBase Shape NoUniqueness), 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
          [TypeBase Shape NoUniqueness]
    ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    mergeinit
          [TypeParamBase VName]
    -> PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([FParam SOACS]
        -> [FParam SOACS]
        -> InternaliseM
             (Body SOACS,
              (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a.
    [TypeParamBase VName]
    -> PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
    -> InternaliseM a
    bindingLoopParams [TypeParamBase VName]
    sparams' PatBase Info VName ParamType
    mergepat [TypeBase Shape NoUniqueness]
    ts (([FParam SOACS]
      -> [FParam SOACS]
      -> InternaliseM
           (Body SOACS,
            (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
     -> InternaliseM
          (Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> ([FParam SOACS]
        -> [FParam 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
    $ \[FParam SOACS]
    shapepat [FParam SOACS]
    mergepat' -> do
            [TypeBase Shape NoUniqueness]
    mergeinit_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    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.
            [SubExp]
    shapeinit <- [VName]
    -> [FParam SOACS]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [SubExp]
    argShapes ((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) [FParam SOACS]
    mergepat' [TypeBase Shape NoUniqueness]
    mergeinit_ts

            (SubExp
    loop_initial_cond, Stms SOACS
    init_loop_cond_stms) <- InternaliseM SubExp
    -> InternaliseM (SubExp, 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 SubExp
     -> InternaliseM (SubExp, Stms (Rep InternaliseM)))
    -> InternaliseM SubExp
    -> InternaliseM (SubExp, Stms (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $ do
              [(Param DeclType, SubExp)]
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param DeclType]
    [FParam SOACS]
    shapepat [SubExp]
    shapeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ \(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
              [(Param DeclType, SubExp)]
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param DeclType]
    [FParam SOACS]
    mergepat' [SubExp]
    mergeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ \(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 NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType (TypeBase Shape NoUniqueness -> Bool)
    -> TypeBase Shape NoUniqueness -> Bool
    forall a b. (a -> b) -> a -> b
    $ Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    paramType Param DeclType
    p ->
                            [SubExp] -> VName -> Exp SOACS
    forall rep. [SubExp] -> VName -> Exp rep
    shapeCoerce (TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
    -> TypeBase Shape NoUniqueness -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    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).
              (Stms SOACS
    cond_stms, SubExp
    cond') <-
                (SubExp -> Stms SOACS -> InternaliseM (Stms SOACS, SubExp))
    -> (SubExp, Stms SOACS) -> InternaliseM (Stms SOACS, SubExp)
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ((Stms SOACS -> SubExp -> InternaliseM (Stms SOACS, SubExp))
    -> SubExp -> Stms SOACS -> InternaliseM (Stms SOACS, SubExp)
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Stms SOACS -> SubExp -> InternaliseM (Stms SOACS, SubExp)
    forall (m :: * -> *) rep a.
    (MonadFreshNames m, Renameable rep, Rename a) =>
    Stms rep -> a -> m (Stms rep, a)
    renameStmsWith)
                  ((SubExp, Stms SOACS) -> InternaliseM (Stms SOACS, SubExp))
    -> InternaliseM (SubExp, Stms SOACS)
    -> InternaliseM (Stms SOACS, SubExp)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< InternaliseM SubExp
    -> InternaliseM (SubExp, 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 ([Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "loop_cond" Exp
    cond)
              Stms (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep InternaliseM)
    Stms SOACS
    cond_stms
              SubExp -> InternaliseM SubExp
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    cond'

            Stms (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep InternaliseM)
    Stms SOACS
    init_loop_cond_stms

            InternaliseM
      (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
    bodyFromStms (InternaliseM
       (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
     -> InternaliseM
          (Body SOACS,
           (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
    -> InternaliseM
         (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a b. (a -> b) -> a -> b
    $ do
              [SubExp]
    ses <- InternaliseM [SubExp]
    loopBody
              [TypeBase Shape NoUniqueness]
    sets <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    ses
              Param DeclType
    loop_while <- [Char] -> DeclType -> InternaliseM (Param DeclType)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "loop_while" (DeclType -> InternaliseM (Param DeclType))
    -> DeclType -> InternaliseM (Param DeclType)
    forall a b. (a -> b) -> a -> b
    $ PrimType -> DeclType
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    I.Bool
              [SubExp]
    shapeargs <- [VName]
    -> [FParam SOACS]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [SubExp]
    argShapes ((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) [FParam SOACS]
    mergepat' [TypeBase Shape NoUniqueness]
    sets

              -- Careful not to clobber anything.
              Body SOACS
    loop_end_cond_body <- Body SOACS -> InternaliseM (Body SOACS)
    forall rep (m :: * -> *).
    (Renameable rep, MonadFreshNames m) =>
    Body rep -> m (Body rep)
    renameBody (Body SOACS -> InternaliseM (Body SOACS))
    -> (InternaliseM Result -> InternaliseM (Body SOACS))
    -> InternaliseM Result
    -> InternaliseM (Body SOACS)
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
    InternaliseM Result -> InternaliseM (Body SOACS)
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (InternaliseM Result -> InternaliseM (Body SOACS))
    -> InternaliseM Result -> InternaliseM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $ do
                [(Param DeclType, SubExp)]
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param DeclType]
    [FParam SOACS]
    shapepat [SubExp]
    shapeargs) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ \(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
                [(Param DeclType, SubExp)]
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param DeclType]
    [FParam SOACS]
    mergepat' [SubExp]
    ses) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
    -> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ \(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 NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType (TypeBase Shape NoUniqueness -> Bool)
    -> TypeBase Shape NoUniqueness -> Bool
    forall a b. (a -> b) -> a -> b
    $ Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    paramType Param DeclType
    p ->
                              [SubExp] -> VName -> Exp SOACS
    forall rep. [SubExp] -> VName -> Exp rep
    shapeCoerce (TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
    -> TypeBase Shape NoUniqueness -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    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
                [SubExp] -> Result
    subExpsRes ([SubExp] -> Result)
    -> InternaliseM [SubExp] -> InternaliseM Result
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "loop_cond" Exp
    cond
              Result
    loop_end_cond <- Body (Rep InternaliseM) -> InternaliseM Result
    forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
    bodyBind Body (Rep InternaliseM)
    Body SOACS
    loop_end_cond_body

              (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                ( [SubExp] -> Result
    subExpsRes [SubExp]
    shapeargs Result -> Result -> Result
    forall a. [a] -> [a] -> [a]
    ++ Result
    loop_end_cond Result -> Result -> Result
    forall a. [a] -> [a] -> [a]
    ++ [SubExp] -> Result
    subExpsRes [SubExp]
    ses,
                  ( VName -> LoopForm
    I.WhileLoop (VName -> LoopForm) -> VName -> LoopForm
    forall a b. (a -> b) -> a -> b
    $ Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName Param DeclType
    loop_while,
                    [Param DeclType]
    [FParam SOACS]
    shapepat,
                    Param DeclType
    loop_while Param DeclType -> [Param DeclType] -> [Param DeclType]
    forall a. a -> [a] -> [a]
    : [Param DeclType]
    [FParam SOACS]
    mergepat',
                    SubExp
    loop_initial_cond SubExp -> [SubExp] -> [SubExp]
    forall a. a -> [a] -> [a]
    : [SubExp]
    mergeinit
                  )
                )
    internaliseAppExp [Char]
    desc AppRes
    _ (E.LetWith IdentBase Info VName StructType
    name IdentBase Info VName StructType
    src SliceBase Info VName
    idxs Exp
    ve Exp
    body SrcLoc
    loc) = do
      let pat :: PatBase Info VName StructType
    pat = VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    E.Id (IdentBase Info VName StructType -> VName
    forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
    E.identName IdentBase Info VName StructType
    name) (IdentBase Info VName StructType -> Info StructType
    forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
    E.identType IdentBase Info VName StructType
    name) SrcLoc
    loc
          src_t :: Info StructType
    src_t = IdentBase Info VName StructType -> Info StructType
    forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
    E.identType IdentBase Info VName StructType
    src
          e :: Exp
    e = Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
    E.Update (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 -> QualName VName) -> VName -> QualName VName
    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
    src) Info StructType
    src_t SrcLoc
    loc) SliceBase Info VName
    idxs Exp
    ve SrcLoc
    loc
      [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    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
          ([SizeBinder VName]
    -> PatBase Info VName StructType -> Exp -> Exp -> SrcLoc -> AppExp
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    E.LetPat [] PatBase Info VName StructType
    pat Exp
    e Exp
    body SrcLoc
    loc)
          (AppRes -> Info AppRes
    forall a. a -> Info a
    Info (StructType -> [VName] -> AppRes
    AppRes (Exp -> StructType
    E.typeOf Exp
    body) [VName]
    forall a. Monoid a => a
    mempty))
    internaliseAppExp [Char]
    desc AppRes
    _ (E.Match Exp
    e NonEmpty (CaseBase Info VName)
    orig_cs SrcLoc
    _) = do
      [SubExp]
    ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_scrutinee") Exp
    e
      NonEmpty (Case (InternaliseM (Body SOACS)))
    cs <- (CaseBase Info VName
     -> InternaliseM (Case (InternaliseM (Body SOACS))))
    -> NonEmpty (CaseBase Info VName)
    -> InternaliseM (NonEmpty (Case (InternaliseM (Body SOACS))))
    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) -> NonEmpty a -> m (NonEmpty b)
    mapM ([SubExp]
    -> CaseBase Info VName
    -> InternaliseM (Case (InternaliseM (Body SOACS)))
    onCase [SubExp]
    ses) NonEmpty (CaseBase Info VName)
    orig_cs
      case NonEmpty (Case (InternaliseM (Body SOACS)))
    -> (Case (InternaliseM (Body SOACS)),
        Maybe (NonEmpty (Case (InternaliseM (Body SOACS)))))
    forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
    NE.uncons NonEmpty (Case (InternaliseM (Body SOACS)))
    cs of
        (I.Case [Maybe PrimValue]
    _ InternaliseM (Body SOACS)
    body, Maybe (NonEmpty (Case (InternaliseM (Body SOACS))))
    Nothing) ->
          (Result -> [SubExp])
    -> InternaliseM Result -> 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) -> Result -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp) (InternaliseM Result -> InternaliseM [SubExp])
    -> InternaliseM Result -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Body (Rep InternaliseM) -> InternaliseM Result
    Body SOACS -> InternaliseM Result
    forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
    bodyBind (Body SOACS -> InternaliseM Result)
    -> InternaliseM (Body SOACS) -> InternaliseM Result
    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
          [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    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 (Body (Rep InternaliseM)))]
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    [SubExp]
    -> [Case (m (Body (Rep m)))] -> m (Body (Rep m)) -> m (Exp (Rep m))
    eMatch [SubExp]
    ses (NonEmpty (Case (InternaliseM (Body SOACS)))
    -> [Case (InternaliseM (Body SOACS))]
    forall a. NonEmpty a -> [a]
    NE.init NonEmpty (Case (InternaliseM (Body SOACS)))
    cs) (Case (InternaliseM (Body (Rep InternaliseM)))
    -> InternaliseM (Body (Rep InternaliseM))
    forall body. Case body -> body
    I.caseBody (Case (InternaliseM (Body (Rep InternaliseM)))
     -> InternaliseM (Body (Rep InternaliseM)))
    -> Case (InternaliseM (Body (Rep InternaliseM)))
    -> InternaliseM (Body (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $ NonEmpty (Case (InternaliseM (Body SOACS)))
    -> Case (InternaliseM (Body SOACS))
    forall a. NonEmpty a -> a
    NE.last 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
          ([Maybe PrimValue]
    cmps, [SubExp]
    pertinent) <- PatBase Info VName StructType
    -> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
    generateCond PatBase Info VName StructType
    p [SubExp]
    ses
          Case (InternaliseM (Body SOACS))
    -> InternaliseM (Case (InternaliseM (Body SOACS)))
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Case (InternaliseM (Body SOACS))
     -> InternaliseM (Case (InternaliseM (Body SOACS))))
    -> (InternaliseM (Body SOACS) -> Case (InternaliseM (Body SOACS)))
    -> InternaliseM (Body SOACS)
    -> InternaliseM (Case (InternaliseM (Body SOACS)))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Maybe PrimValue]
    -> InternaliseM (Body SOACS) -> Case (InternaliseM (Body SOACS))
    forall body. [Maybe PrimValue] -> body -> Case body
    I.Case [Maybe PrimValue]
    cmps (InternaliseM (Body SOACS)
     -> InternaliseM (Case (InternaliseM (Body SOACS))))
    -> InternaliseM (Body SOACS)
    -> InternaliseM (Case (InternaliseM (Body SOACS)))
    forall a b. (a -> b) -> a -> b
    $
            [SizeBinder VName]
    -> PatBase Info VName StructType
    -> [SubExp]
    -> InternaliseM (Body SOACS)
    -> InternaliseM (Body SOACS)
    forall a.
    [SizeBinder VName]
    -> PatBase Info VName StructType
    -> [SubExp]
    -> InternaliseM a
    -> InternaliseM a
    internalisePat' [] PatBase Info VName StructType
    p [SubExp]
    pertinent (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
    -> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $
              [Char] -> Exp -> InternaliseM (Body SOACS)
    internaliseBody [Char]
    "case" Exp
    case_e
    internaliseAppExp [Char]
    desc AppRes
    _ (E.If Exp
    ce Exp
    te Exp
    fe SrcLoc
    _) =
      [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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
          (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
    <$> [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "cond" Exp
    ce)
          ([Char] -> Exp -> InternaliseM (Body SOACS)
    internaliseBody ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "_t") Exp
    te)
          ([Char] -> Exp -> InternaliseM (Body SOACS)
    internaliseBody ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "_f") Exp
    fe)
    internaliseAppExp [Char]
    _ 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

    internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
    internaliseExp :: [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc (E.Parens Exp
    e SrcLoc
    _) =
      [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
    internaliseExp [Char]
    desc (E.Hole (Info StructType
    t) SrcLoc
    loc) = do
      let msg :: Text
    msg = Doc Any -> Text
    forall a. Doc a -> Text
    docText (Doc Any -> Text) -> Doc Any -> Text
    forall a b. (a -> b) -> a -> b
    $ Doc Any
    "Reached hole of type: " Doc Any -> Doc Any -> Doc Any
    forall a. Semigroup a => a -> a -> a
    <> Doc Any -> Doc Any
    forall ann. Doc ann -> Doc ann
    align (StructType -> Doc Any
    forall a ann. Pretty a => a -> Doc ann
    forall ann. StructType -> Doc ann
    pretty StructType
    t)
          ts :: [TypeBase ExtShape Uniqueness]
    ts = (Tree (TypeBase ExtShape Uniqueness)
     -> [TypeBase ExtShape Uniqueness])
    -> [Tree (TypeBase ExtShape Uniqueness)]
    -> [TypeBase ExtShape Uniqueness]
    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 Uniqueness)
    -> [TypeBase ExtShape Uniqueness]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList ([Tree (TypeBase ExtShape Uniqueness)]
     -> [TypeBase ExtShape Uniqueness])
    -> [Tree (TypeBase ExtShape Uniqueness)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> a -> b
    $ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
    internaliseType (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    E.toStruct StructType
    t)
      Certs
    c <- [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "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 (TypeBase ExtShape Uniqueness -> Maybe DeclType)
    -> [TypeBase ExtShape Uniqueness] -> Maybe [DeclType]
    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 ExtShape Uniqueness -> Maybe DeclType
    forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
    hasStaticShape [TypeBase ExtShape Uniqueness]
    ts of
        Maybe [DeclType]
    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 Uniqueness] -> [Char]
    forall a. Show a => a -> [Char]
    show [TypeBase ExtShape Uniqueness]
    ts
        Just [DeclType]
    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
    $ (DeclType -> InternaliseM SubExp)
    -> [DeclType] -> 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
    . [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    desc (Exp SOACS -> InternaliseM SubExp)
    -> (DeclType -> InternaliseM (Exp SOACS))
    -> DeclType
    -> InternaliseM SubExp
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< TypeBase Shape NoUniqueness
    -> InternaliseM (Exp (Rep InternaliseM))
    TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
    forall (m :: * -> *).
    MonadBuilder m =>
    TypeBase Shape NoUniqueness -> m (Exp (Rep m))
    eBlank (TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
    -> (DeclType -> TypeBase Shape NoUniqueness)
    -> DeclType
    -> InternaliseM (Exp SOACS)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DeclType -> TypeBase Shape NoUniqueness
    forall shape.
    TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
    I.fromDecl) [DeclType]
    ts'
    internaliseExp [Char]
    desc (E.QualParens (QualName VName, SrcLoc)
    _ Exp
    e SrcLoc
    _) =
      [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
    internaliseExp [Char]
    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
    . [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 NoUniqueness -> 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 NoUniqueness -> BasicOp)
    -> TypeBase Shape NoUniqueness -> BasicOp
    forall a b. (a -> b) -> a -> b
    $
            PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int8
    internaliseExp [Char]
    _ (E.Var (E.QualName [VName]
    _ VName
    name) Info StructType
    _ SrcLoc
    _) = do
      Maybe [SubExp]
    subst <- VName -> InternaliseM (Maybe [SubExp])
    lookupSubst VName
    name
      case Maybe [SubExp]
    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 [Char]
    desc (E.AppExp AppExp
    e (Info AppRes
    appres)) = do
      [SubExp]
    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
    $ [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
    internaliseAppExp [Char]
    desc AppRes
    appres AppExp
    e
      AppRes -> [SubExp] -> InternaliseM ()
    bindExtSizes AppRes
    appres [SubExp]
    ses
      [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp]
    ses
    internaliseExp [Char]
    _ (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 [Char]
    _ (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 [Char]
    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 ([Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc) [Exp]
    es
    internaliseExp [Char]
    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
    <$> [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    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 [Char]
    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
    . [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 [Char]
    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] -> 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 u. Int -> TypeBase dim u -> Maybe (TypeBase dim u)
    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
          [VName]
    flat_arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "flat_literal" Exp
    flat_lit
          [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    flat_arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
    -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    flat_arr -> do
            TypeBase Shape NoUniqueness
    flat_arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    flat_arr
            let new_shape' :: Shape
    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 NoUniqueness -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    I.arrayShape TypeBase Shape NoUniqueness
    flat_arr_t
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 -> NewShape SubExp -> BasicOp
    I.Reshape VName
    flat_arr (Shape -> Shape -> NewShape SubExp
    forall old new.
    ArrayShape old =>
    old -> ShapeBase new -> NewShape new
    reshapeAll (TypeBase Shape NoUniqueness -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    I.arrayShape TypeBase Shape NoUniqueness
    flat_arr_t) Shape
    new_shape')
      | Bool
    otherwise = do
          [[SubExp]]
    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 ([Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "arr_elem") [Exp]
    es
          let arr_t_ext :: [TypeBase ExtShape Uniqueness]
    arr_t_ext = (Tree (TypeBase ExtShape Uniqueness)
     -> [TypeBase ExtShape Uniqueness])
    -> [Tree (TypeBase ExtShape Uniqueness)]
    -> [TypeBase ExtShape Uniqueness]
    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 Uniqueness)
    -> [TypeBase ExtShape Uniqueness]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList ([Tree (TypeBase ExtShape Uniqueness)]
     -> [TypeBase ExtShape Uniqueness])
    -> [Tree (TypeBase ExtShape Uniqueness)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> a -> b
    $ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
    internaliseType (StructType -> [Tree (TypeBase ExtShape Uniqueness)])
    -> StructType -> [Tree (TypeBase ExtShape Uniqueness)]
    forall a b. (a -> b) -> a -> b
    $ StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    E.toStruct StructType
    arr_t

          [TypeBase Shape NoUniqueness]
    rowtypes <-
            case (TypeBase ExtShape Uniqueness
     -> Maybe (TypeBase Shape NoUniqueness))
    -> [TypeBase ExtShape Uniqueness]
    -> Maybe [TypeBase Shape NoUniqueness]
    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 Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> Maybe (TypeBase Shape NoUniqueness)
    -> Maybe (TypeBase Shape NoUniqueness)
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType (Maybe (TypeBase Shape NoUniqueness)
     -> Maybe (TypeBase Shape NoUniqueness))
    -> (TypeBase ExtShape Uniqueness
        -> Maybe (TypeBase Shape NoUniqueness))
    -> TypeBase ExtShape Uniqueness
    -> Maybe (TypeBase Shape NoUniqueness)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ExtType -> Maybe (TypeBase Shape NoUniqueness)
    forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
    hasStaticShape (ExtType -> Maybe (TypeBase Shape NoUniqueness))
    -> (TypeBase ExtShape Uniqueness -> ExtType)
    -> TypeBase ExtShape Uniqueness
    -> Maybe (TypeBase Shape NoUniqueness)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase ExtShape Uniqueness -> ExtType
    forall shape.
    TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
    I.fromDecl) [TypeBase ExtShape Uniqueness]
    arr_t_ext of
              Just [TypeBase Shape NoUniqueness]
    ts -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [TypeBase Shape NoUniqueness]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [TypeBase Shape NoUniqueness]
    ts
              Maybe [TypeBase Shape NoUniqueness]
    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 NoUniqueness]
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> InternaliseM [TypeBase Shape NoUniqueness])
    -> [Char] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    e'

          let arraylit :: [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp rep)
    arraylit [SubExp]
    ks TypeBase Shape NoUniqueness
    rt = do
                [SubExp]
    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 NoUniqueness
    -> [Char]
    -> SubExp
    -> InternaliseM SubExp
    ensureShape
                        ErrorMsg SubExp
    "shape of element differs from shape of first element"
                        TypeBase Shape NoUniqueness
    rt
                        [Char]
    "elem_reshaped"
                    )
                    [SubExp]
    ks
                Exp rep -> InternaliseM (Exp rep)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp rep -> InternaliseM (Exp rep))
    -> Exp rep -> InternaliseM (Exp rep)
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp rep
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
    I.ArrayLit [SubExp]
    ks' TypeBase Shape NoUniqueness
    rt

          (Exp SOACS -> InternaliseM SubExp)
    -> [Exp SOACS] -> 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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    desc)
            ([Exp SOACS] -> InternaliseM [SubExp])
    -> InternaliseM [Exp SOACS] -> InternaliseM [SubExp]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< if [[SubExp]] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [[SubExp]]
    es'
              then (TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
    -> [TypeBase Shape NoUniqueness] -> InternaliseM [Exp SOACS]
    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] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
    forall {rep}.
    [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp rep)
    arraylit []) [TypeBase Shape NoUniqueness]
    rowtypes
              else ([SubExp]
     -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
    -> [[SubExp]]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [Exp SOACS]
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m [c]
    zipWithM [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
    forall {rep}.
    [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp rep)
    arraylit ([[SubExp]] -> [[SubExp]]
    forall a. [[a]] -> [[a]]
    transpose [[SubExp]]
    es') [TypeBase Shape NoUniqueness]
    rowtypes
      where
        isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
        isArrayLiteral :: Exp -> Maybe ([Int], [Exp])
    isArrayLiteral (E.ArrayLit [Exp]
    inner_es Info StructType
    _ SrcLoc
    _) = do
          ([Int]
    eshape, [Exp]
    e) : [([Int], [Exp])]
    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
          Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ (([Int], [Exp]) -> Bool) -> [([Int], [Exp])] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (([Int]
    eshape ==) ([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])]
    inner_es'
          ([Int], [Exp]) -> Maybe ([Int], [Exp])
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Exp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Exp]
    inner_es Int -> [Int] -> [Int]
    forall a. a -> [a] -> [a]
    : [Int]
    eshape, [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])]
    inner_es')
        isArrayLiteral Exp
    e =
          ([Int], [Exp]) -> Maybe ([Int], [Exp])
    forall a. a -> Maybe a
    Just ([], [Exp
    e])
    internaliseExp [Char]
    desc (E.Ascript Exp
    e TypeExp Exp VName
    _ SrcLoc
    _) =
      [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
    internaliseExp [Char]
    desc (E.Coerce Exp
    e TypeExp Exp VName
    _ (Info StructType
    et) SrcLoc
    _) = do
      [SubExp]
    ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
      [TypeBase ExtShape Uniqueness]
    ts <- StructType
    -> [TypeBase Shape NoUniqueness] -> [TypeBase ExtShape Uniqueness]
    forall shape u.
    StructType -> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
    internaliseCoerceType (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    E.toStruct StructType
    et) ([TypeBase Shape NoUniqueness] -> [TypeBase ExtShape Uniqueness])
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM [TypeBase ExtShape Uniqueness]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    ses
      [ErrorMsgPart SubExp]
    dt' <- StructType -> InternaliseM [ErrorMsgPart SubExp]
    forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
    typeExpForError (StructType -> InternaliseM [ErrorMsgPart SubExp])
    -> StructType -> InternaliseM [ErrorMsgPart SubExp]
    forall a b. (a -> b) -> a -> b
    $ StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct StructType
    et
      [(SubExp, TypeBase ExtShape Uniqueness)]
    -> ((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([SubExp]
    -> [TypeBase ExtShape Uniqueness]
    -> [(SubExp, TypeBase ExtShape Uniqueness)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    ses [TypeBase ExtShape Uniqueness]
    ts) (((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
     -> InternaliseM [SubExp])
    -> ((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(SubExp
    e', TypeBase ExtShape Uniqueness
    t') -> do
        [SubExp]
    dims <- TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType SubExp
    e'
        let parts :: [ErrorMsgPart SubExp]
    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
    "\"."]
        ErrorMsg SubExp
    -> ExtType -> [Char] -> SubExp -> InternaliseM SubExp
    ensureExtShape ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
    forall a. [ErrorMsgPart a] -> ErrorMsg a
    errorMsg [ErrorMsgPart SubExp]
    parts) (TypeBase ExtShape Uniqueness -> ExtType
    forall shape.
    TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
    I.fromDecl TypeBase ExtShape Uniqueness
    t') [Char]
    desc SubExp
    e'
    internaliseExp [Char]
    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
      SubExp
    e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "negate_arg" Exp
    e
      TypeBase Shape NoUniqueness
    et <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType SubExp
    e'
      case TypeBase Shape NoUniqueness
    et of
        I.Prim PrimType
    pt ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 NoUniqueness
    _ -> [Char] -> InternaliseM [SubExp]
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Futhark.Internalise.internaliseExp: non-primitive type in Negate"
    internaliseExp [Char]
    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
      SubExp
    e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "not_arg" Exp
    e
      TypeBase Shape NoUniqueness
    et <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType SubExp
    e'
      case TypeBase Shape NoUniqueness
    et of
        I.Prim (I.IntType IntType
    t) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 NoUniqueness
    _ ->
          [Char] -> InternaliseM [SubExp]
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Futhark.Internalise.internaliseExp: non-int/bool type in Not"
    internaliseExp [Char]
    desc (E.Update Exp
    src SliceBase Info VName
    slice Exp
    ve 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
      [SubExp]
    ves <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "lw_val" Exp
    ve
      [VName]
    srcs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "src" Exp
    src
      ([SubExp]
    src_dims, [SubExp]
    ve_dims) <- case ([VName]
    srcs, [SubExp]
    ves) of
        (VName
    src_v : [VName]
    _, SubExp
    ve_v : [SubExp]
    _) ->
          (,)
            ([SubExp] -> [SubExp] -> ([SubExp], [SubExp]))
    -> InternaliseM [SubExp]
    -> InternaliseM ([SubExp] -> ([SubExp], [SubExp]))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    src_v)
            InternaliseM ([SubExp] -> ([SubExp], [SubExp]))
    -> InternaliseM [SubExp] -> InternaliseM ([SubExp], [SubExp])
    forall a b.
    InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType SubExp
    ve_v)
        ([VName], [SubExp])
    _ -> ([SubExp], [SubExp]) -> InternaliseM ([SubExp], [SubExp])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([], []) -- Will this happen?
      ([DimIndex SubExp]
    idxs', Certs
    cs) <- [SubExp]
    -> SliceBase Info VName -> InternaliseM ([DimIndex SubExp], Certs)
    internaliseSlice [SubExp]
    src_dims SliceBase Info VName
    slice
      let src_dims' :: [SubExp]
    src_dims' = Slice SubExp -> [SubExp]
    forall d. Slice d -> [d]
    sliceDims ([DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [DimIndex SubExp]
    idxs')
          rank :: Int
    rank = [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    src_dims'
          errormsg :: ErrorMsg SubExp
    errormsg =
            ErrorMsg SubExp
    "Shape "
              ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
    forall a. Semigroup a => a -> a -> a
    <> [SubExp] -> ErrorMsg SubExp
    forall a. [a] -> ErrorMsg a
    errorShape [SubExp]
    src_dims'
              ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
    forall a. Semigroup a => a -> a -> a
    <> ErrorMsg SubExp
    " of slice does not match shape "
              ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
    forall a. Semigroup a => a -> a -> a
    <> [SubExp] -> ErrorMsg SubExp
    forall a. [a] -> ErrorMsg a
    errorShape (Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take Int
    rank [SubExp]
    ve_dims)
              ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
    forall a. Semigroup a => a -> a -> a
    <> ErrorMsg SubExp
    " of value."

      let comb :: VName -> SubExp -> InternaliseM VName
    comb VName
    sname SubExp
    ve' = do
            TypeBase Shape NoUniqueness
    sname_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    sname
            let full_slice :: Slice SubExp
    full_slice = TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
    fullSlice TypeBase Shape NoUniqueness
    sname_t [DimIndex SubExp]
    idxs'
                rowtype :: TypeBase Shape NoUniqueness
    rowtype = TypeBase Shape NoUniqueness
    sname_t TypeBase Shape NoUniqueness
    -> [SubExp] -> TypeBase Shape NoUniqueness
    forall oldshape u.
    TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
    `setArrayDims` Slice SubExp -> [SubExp]
    forall d. Slice d -> [d]
    sliceDims Slice SubExp
    full_slice
            SubExp
    ve'' <-
              ErrorMsg SubExp
    -> TypeBase Shape NoUniqueness
    -> [Char]
    -> SubExp
    -> InternaliseM SubExp
    ensureShape ErrorMsg SubExp
    errormsg TypeBase Shape NoUniqueness
    rowtype [Char]
    "lw_val_correct_shape" SubExp
    ve'
            [Char]
    -> VName
    -> Slice SubExp
    -> Exp (Rep InternaliseM)
    -> InternaliseM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
    letInPlace [Char]
    desc VName
    sname Slice SubExp
    full_slice (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
    BasicOp (BasicOp -> Exp (Rep InternaliseM))
    -> BasicOp -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp SubExp
    ve''
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a. Certs -> InternaliseM a -> InternaliseM a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying Certs
    cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ (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
    <$> (VName -> SubExp -> InternaliseM VName)
    -> [VName] -> [SubExp] -> InternaliseM [VName]
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m [c]
    zipWithM VName -> SubExp -> InternaliseM VName
    comb [VName]
    srcs [SubExp]
    ves
    internaliseExp [Char]
    desc (E.RecordUpdate Exp
    src [Name]
    fields 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
      [SubExp]
    src' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    src
      [SubExp]
    ve' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    ve
      StructType
    -> [Name] -> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
    forall {m :: * -> *} {als} {a}.
    Monad m =>
    TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
    replace (Exp -> StructType
    E.typeOf Exp
    src) [Name]
    fields [SubExp]
    ve' [SubExp]
    src'
      where
        replace :: TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
    replace (E.Scalar (E.Record Map Name (TypeBase Exp als)
    m)) (Name
    f : [Name]
    fs) [a]
    ve' [a]
    src'
          | Just TypeBase Exp als
    t <- Name -> Map Name (TypeBase Exp als) -> Maybe (TypeBase Exp als)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Name
    f Map Name (TypeBase Exp als)
    m = 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)
    -> ([(Name, TypeBase Exp als)] -> [Int])
    -> [(Name, TypeBase Exp als)]
    -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((Name, TypeBase Exp als) -> Int)
    -> [(Name, TypeBase Exp als)] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map (TypeBase Exp als -> Int
    forall als. TypeBase Exp als -> Int
    internalisedTypeSize (TypeBase Exp als -> Int)
    -> ((Name, TypeBase Exp als) -> TypeBase Exp als)
    -> (Name, TypeBase Exp als)
    -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Name, TypeBase Exp als) -> TypeBase Exp als
    forall a b. (a, b) -> b
    snd) ([(Name, TypeBase Exp als)] -> Int)
    -> [(Name, TypeBase Exp als)] -> Int
    forall a b. (a -> b) -> a -> b
    $
                      ((Name, TypeBase Exp als) -> Bool)
    -> [(Name, TypeBase Exp als)] -> [(Name, TypeBase Exp als)]
    forall a. (a -> Bool) -> [a] -> [a]
    takeWhile ((Name -> Name -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Name
    f) (Name -> Bool)
    -> ((Name, TypeBase Exp als) -> Name)
    -> (Name, TypeBase Exp als)
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Name, TypeBase Exp als) -> Name
    forall a b. (a, b) -> a
    fst) ([(Name, TypeBase Exp als)] -> [(Name, TypeBase Exp als)])
    -> (Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)])
    -> Map Name (TypeBase Exp als)
    -> [(Name, TypeBase Exp als)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)]
    forall a. Map Name a -> [(Name, a)]
    sortFields (Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)])
    -> Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)]
    forall a b. (a -> b) -> a -> b
    $
                        Map Name (TypeBase Exp als)
    m
                  k :: Int
    k = TypeBase Exp als -> Int
    forall als. TypeBase Exp als -> Int
    internalisedTypeSize TypeBase Exp als
    t
                  ([a]
    bef, [a]
    to_update, [a]
    aft) = Int -> Int -> [a] -> ([a], [a], [a])
    forall a. Int -> Int -> [a] -> ([a], [a], [a])
    splitAt3 Int
    i Int
    k [a]
    src'
              [a]
    src'' <- TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
    replace TypeBase Exp als
    t [Name]
    fs [a]
    ve' [a]
    to_update
              [a] -> m [a]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([a] -> m [a]) -> [a] -> m [a]
    forall a b. (a -> b) -> a -> b
    $ [a]
    bef [a] -> [a] -> [a]
    forall a. [a] -> [a] -> [a]
    ++ [a]
    src'' [a] -> [a] -> [a]
    forall a. [a] -> [a] -> [a]
    ++ [a]
    aft
        replace TypeBase Exp als
    _ [Name]
    _ [a]
    ve' [a]
    _ = [a] -> m [a]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [a]
    ve'
    internaliseExp [Char]
    desc (E.Attr AttrInfo VName
    attr Exp
    e SrcLoc
    loc) = do
      Attr
    attr' <- AttrInfo VName -> InternaliseM Attr
    internaliseAttr AttrInfo VName
    attr
      [SubExp]
    e' <- (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 (Attr -> InternaliseEnv -> InternaliseEnv
    f Attr
    attr') (InternaliseM [SubExp] -> InternaliseM [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
      case Attr
    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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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
          [TypeBase Shape NoUniqueness]
    ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    e'
          [(TypeBase Shape NoUniqueness, SubExp)]
    -> ((TypeBase Shape NoUniqueness, SubExp) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([TypeBase Shape NoUniqueness]
    -> [SubExp] -> [(TypeBase Shape NoUniqueness, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [TypeBase Shape NoUniqueness]
    ts [SubExp]
    e') (((TypeBase Shape NoUniqueness, SubExp) -> InternaliseM SubExp)
     -> InternaliseM [SubExp])
    -> ((TypeBase Shape NoUniqueness, SubExp) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(TypeBase Shape NoUniqueness
    t, SubExp
    se) ->
            case TypeBase Shape NoUniqueness
    t of
              I.Array PrimType
    pt Shape
    shape NoUniqueness
    _ ->
                [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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
              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 NoUniqueness
    _ -> SubExp -> InternaliseM SubExp
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    se
        Attr
    "blank" -> do
          [TypeBase Shape NoUniqueness]
    ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    e'
          [(TypeBase Shape NoUniqueness, SubExp)]
    -> ((TypeBase Shape NoUniqueness, SubExp) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([TypeBase Shape NoUniqueness]
    -> [SubExp] -> [(TypeBase Shape NoUniqueness, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [TypeBase Shape NoUniqueness]
    ts [SubExp]
    e') (((TypeBase Shape NoUniqueness, SubExp) -> InternaliseM SubExp)
     -> InternaliseM [SubExp])
    -> ((TypeBase Shape NoUniqueness, SubExp) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(TypeBase Shape NoUniqueness
    t, SubExp
    se) ->
            case TypeBase Shape NoUniqueness
    t of
              I.Array PrimType
    pt Shape
    shape NoUniqueness
    _ ->
                [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 NoUniqueness
    _ -> SubExp -> InternaliseM SubExp
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    se
        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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 [Char]
    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
      SubExp
    e1' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "assert_cond" Exp
    e1
      Certs
    c <- [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "assert_c" SubExp
    e1' (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 -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
    forall a b. (a -> b) -> a -> b
    $ Text
    "Assertion is false: " Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    check]
      -- Make sure there are some bindings to certify.
      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
    $ (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 SubExp -> InternaliseM SubExp
    forall {m :: * -> *}. MonadBuilder m => SubExp -> m SubExp
    rebind ([SubExp] -> InternaliseM [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e2
      where
        rebind :: SubExp -> m SubExp
    rebind SubExp
    v = do
          VName
    v' <- [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "assert_res"
          [VName] -> Exp (Rep m) -> m ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    v'] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
    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
    $ SubExp -> BasicOp
    I.SubExp SubExp
    v
          SubExp -> m SubExp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp -> m SubExp) -> SubExp -> m SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    I.Var VName
    v'
    internaliseExp [Char]
    _ (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
      ([TypeBase ExtShape Uniqueness]
    ts, [(Name, [Int])]
    constr_map) <- Map Name [StructType]
    -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
    internaliseSumType (Map Name [StructType]
     -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])]))
    -> Map Name [StructType]
    -> InternaliseM ([TypeBase ExtShape Uniqueness], [(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 u. TypeBase dim u -> TypeBase dim NoUniqueness
    E.toStruct) Map Name [StructType]
    fs
      [SubExp]
    es' <- [[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 ([Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "payload") [Exp]
    es

      let noExt :: p -> f SubExp
    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
      [TypeBase Shape NoUniqueness]
    ts' <- (Int -> InternaliseM SubExp)
    -> [ExtType] -> InternaliseM [TypeBase Shape NoUniqueness]
    forall (m :: * -> *) u.
    Monad m =>
    (Int -> m SubExp) -> [TypeBase ExtShape u] -> m [TypeBase Shape u]
    instantiateShapes Int -> InternaliseM SubExp
    forall {f :: * -> *} {p}. Applicative f => p -> f SubExp
    noExt ([ExtType] -> InternaliseM [TypeBase Shape NoUniqueness])
    -> [ExtType] -> InternaliseM [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ (TypeBase ExtShape Uniqueness -> ExtType)
    -> [TypeBase ExtShape Uniqueness] -> [ExtType]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase ExtShape Uniqueness -> ExtType
    forall shape.
    TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
    fromDecl [TypeBase ExtShape Uniqueness]
    ts

      case Name -> [(Name, [Int])] -> Maybe (Int, [Int])
    forall a b. Eq a => a -> [(a, b)] -> Maybe (Int, b)
    lookupWithIndex Name
    c [(Name, [Int])]
    constr_map of
        Just (Int
    i, [Int]
    js) ->
          (IntType -> Integer -> SubExp
    intConst IntType
    Int8 (Int -> Integer
    forall a. Integral a => a -> Integer
    toInteger Int
    i) :) ([SubExp] -> [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Int
    -> [TypeBase Shape NoUniqueness]
    -> [(Int, SubExp)]
    -> InternaliseM [SubExp]
    forall {f :: * -> *} {t}.
    (Num t, MonadBuilder f, Eq t) =>
    t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
    clauses Int
    0 [TypeBase Shape NoUniqueness]
    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 NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
    clauses t
    j (TypeBase Shape NoUniqueness
    t : [TypeBase Shape NoUniqueness]
    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]) -> f [SubExp] -> f [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
    clauses (t
    j t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    1) [TypeBase Shape NoUniqueness]
    ts [(t, SubExp)]
    js_to_es
          | Bool
    otherwise = do
              SubExp
    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.
                [Char] -> Exp (Rep f) -> f SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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 NoUniqueness
    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 NoUniqueness -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    I.arrayShape TypeBase Shape NoUniqueness
    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 NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    t
                    TypeBase Shape NoUniqueness
    _ -> TypeBase Shape NoUniqueness -> f (Exp (Rep f))
    forall (m :: * -> *).
    MonadBuilder m =>
    TypeBase Shape NoUniqueness -> m (Exp (Rep m))
    eBlank TypeBase Shape NoUniqueness
    t
              (SubExp
    blank :) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
    clauses (t
    j t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    1) [TypeBase Shape NoUniqueness]
    ts [(t, SubExp)]
    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 [Char]
    _ (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 [Char]
    _ (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 [Char]
    _ (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 [Char]
    _ (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 [Char]
    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 als. TypeBase Exp als -> 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 als. TypeBase Exp als -> 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
    <$> [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
    internaliseExp [Char]
    _ 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 [Char]
    _ 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 [Char]
    _ 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 [Char]
    _ 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 [Char]
    _ e :: Exp
    e@E.ProjectSection {} =
      [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)
    internaliseExp [Char]
    _ e :: Exp
    e@E.IndexSection {} =
      [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 index 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 :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
    internaliseArg :: [Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
    internaliseArg [Char]
    desc (Exp
    arg, Maybe VName
    argdim) = do
      Scope SOACS
    exists <- InternaliseM (Scope SOACS)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
      case Maybe VName
    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
          [SubExp]
    arg' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    arg
          case ([SubExp]
    arg', Maybe VName
    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 ()
          [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp]
    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
      ([Maybe PrimValue]
    cmps, [SubExp]
    pertinent, [SubExp]
    _) <- 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
      ([Maybe PrimValue], [SubExp])
    -> InternaliseM ([Maybe PrimValue], [SubExp])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Maybe PrimValue]
    cmps, [SubExp]
    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
    _ : [a]
    ses) = do
          ([TypeBase ExtShape Uniqueness]
    payload_ts, [(Name, [Int])]
    m) <- Map Name [StructType]
    -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
    internaliseSumType (Map Name [StructType]
     -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])]))
    -> Map Name [StructType]
    -> InternaliseM ([TypeBase ExtShape Uniqueness], [(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 u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct) Map Name [StructType]
    fs
          case Name -> [(Name, [Int])] -> Maybe (Int, [Int])
    forall a b. Eq a => a -> [(a, b)] -> Maybe (Int, b)
    lookupWithIndex Name
    c [(Name, [Int])]
    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 Uniqueness] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TypeBase ExtShape Uniqueness]
    payload_ts) [a]
    ses
              ([Maybe PrimValue]
    cmps, [a]
    pertinent, [a]
    _) <-
                [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 !!) [Int]
    payload_is
              let missingCmps :: Int -> a -> Maybe PrimValue
    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
              ([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 (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
    Int8 (Integer -> IntValue) -> Integer -> IntValue
    forall a b. (a -> b) -> a -> b
    $ Int -> Integer
    forall a. Integral a => a -> Integer
    toInteger Int
    tag)
                    Maybe PrimValue -> [Maybe PrimValue] -> [Maybe PrimValue]
    forall a. a -> [a] -> [a]
    : (Int -> a -> Maybe PrimValue) -> [Int] -> [a] -> [Maybe PrimValue]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Int -> a -> Maybe PrimValue
    missingCmps [Int
    0 ..] [a]
    payload_ses,
                  [a]
    pertinent,
                  [a]
    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 als. TypeBase Exp als -> Int
    internalisedTypeSize (StructType -> Int) -> StructType -> Int
    forall a b. (a -> b) -> a -> b
    $ StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    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 NoUniqueness -> StructType
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
    -> ScalarTypeBase Exp NoUniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ Map Name StructType -> ScalarTypeBase Exp NoUniqueness
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    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 NoUniqueness -> StructType
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
    -> ScalarTypeBase Exp NoUniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ Map Name StructType -> ScalarTypeBase Exp NoUniqueness
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    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
          ([Maybe PrimValue]
    cmps1, [a]
    pertinent1, [a]
    ses') <- PatBase Info vn StructType
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    compares PatBase Info vn StructType
    pat [a]
    ses
          ([Maybe PrimValue]
    cmps2, [a]
    pertinent2, [a]
    ses'') <- [PatBase Info vn StructType]
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    comparesMany [PatBase Info vn StructType]
    pats [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
            ( [Maybe PrimValue]
    cmps1 [Maybe PrimValue] -> [Maybe PrimValue] -> [Maybe PrimValue]
    forall a. Semigroup a => a -> a -> a
    <> [Maybe PrimValue]
    cmps2,
              [a]
    pertinent1 [a] -> [a] -> [a]
    forall a. Semigroup a => a -> a -> a
    <> [a]
    pertinent2,
              [a]
    ses''
            )

    internalisePat ::
      String ->
      [E.SizeBinder VName] ->
      E.Pat StructType ->
      E.Exp ->
      InternaliseM a ->
      InternaliseM a
    internalisePat :: forall a.
    [Char]
    -> [SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> InternaliseM a
    -> InternaliseM a
    internalisePat [Char]
    desc [SizeBinder VName]
    sizes PatBase Info VName StructType
    p Exp
    e InternaliseM a
    m = do
      [SubExp]
    ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc' Exp
    e
      [SizeBinder VName]
    -> PatBase Info VName StructType
    -> [SubExp]
    -> InternaliseM a
    -> InternaliseM a
    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
      where
        desc' :: [Char]
    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 -> [Char]
    baseString (VName -> [Char]) -> VName -> [Char]
    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]
    _ -> [Char]
    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
      [TypeBase Shape NoUniqueness]
    ses_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    ses
      PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([VName] -> InternaliseM a)
    -> InternaliseM a
    forall a.
    PatBase Info VName ParamType
    -> [TypeBase Shape NoUniqueness]
    -> ([VName] -> InternaliseM a)
    -> InternaliseM a
    stmPat (Diet -> StructType -> ParamType
    forall u. Diet -> TypeBase Exp u -> 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
    p) [TypeBase Shape NoUniqueness]
    ses_ts (([VName] -> InternaliseM a) -> InternaliseM a)
    -> ([VName] -> InternaliseM a) -> InternaliseM a
    forall a b. (a -> b) -> a -> b
    $ \[VName]
    pat_names -> do
        AppRes -> [SubExp] -> InternaliseM ()
    bindExtSizes (StructType -> [VName] -> AppRes
    AppRes (PatBase Info VName StructType -> StructType
    forall d u. Pat (TypeBase d u) -> TypeBase d u
    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
      ([DimIndex SubExp]
    idxs', [SubExp]
    oks, [[ErrorMsgPart SubExp]]
    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
      SubExp
    ok <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "index_ok" (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAll [SubExp]
    oks
      let msg :: ErrorMsg SubExp
    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
    "]."]
      Certs
    c <- [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "index_certs" SubExp
    ok ErrorMsg SubExp
    msg
      ([DimIndex SubExp], Certs)
    -> InternaliseM ([DimIndex SubExp], Certs)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([DimIndex SubExp]
    idxs', Certs
    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
      (SubExp
    i', IntType
    _) <- [Char] -> Exp -> InternaliseM (SubExp, IntType)
    internaliseSizeExp [Char]
    "i" Exp
    i
      let lowerBound :: Exp SOACS
    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 :: Exp SOACS
    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
      SubExp
    ok <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "bounds_check" (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< BinOp
    -> InternaliseM (Exp (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eBinOp BinOp
    I.LogAnd (Exp SOACS -> InternaliseM (Exp SOACS)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp SOACS
    lowerBound) (Exp SOACS -> InternaliseM (Exp SOACS)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp SOACS
    upperBound)
      (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    I.DimFix SubExp
    i', SubExp
    ok, [PrimType -> SubExp -> ErrorMsgPart SubExp
    forall a. PrimType -> a -> ErrorMsgPart a
    ErrorVal PrimType
    int64 SubExp
    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
        SubExp
    w_minus_1 <-
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
        (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
          ( SubExp -> SubExp -> SubExp -> DimIndex SubExp
    forall d. d -> d -> d -> DimIndex d
    I.DimSlice SubExp
    w_minus_1 SubExp
    w (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ IntType -> Integer -> SubExp
    intConst IntType
    Int64 (-Integer
    1),
            Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    True,
            [ErrorMsgPart SubExp]
    forall a. Monoid a => a
    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
      SubExp
    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
    . [Char] -> Exp -> InternaliseM (SubExp, IntType)
    internaliseSizeExp [Char]
    "s") Maybe Exp
    s
      SubExp
    s_sign <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "s_sign" (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
    $ UnOp -> SubExp -> BasicOp
    I.UnOp (IntType -> UnOp
    I.SSignum IntType
    Int64) SubExp
    s'
      SubExp
    backwards <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "backwards" (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
    s_sign SubExp
    negone
      SubExp
    w_minus_1 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
      let i_def :: InternaliseM SubExp
    i_def =
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    backwards)
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    w_minus_1])
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    zero])
          j_def :: InternaliseM SubExp
    j_def =
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    backwards)
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    negone])
                ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    w])
      SubExp
    i' <- InternaliseM SubExp
    -> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe InternaliseM SubExp
    i_def (((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
    . [Char] -> Exp -> InternaliseM (SubExp, IntType)
    internaliseSizeExp [Char]
    "i") Maybe Exp
    i
      SubExp
    j' <- InternaliseM SubExp
    -> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe InternaliseM SubExp
    j_def (((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
    . [Char] -> Exp -> InternaliseM (SubExp, IntType)
    internaliseSizeExp [Char]
    "j") Maybe Exp
    j
      SubExp
    j_m_i <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "j_m_i" (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
    j' SubExp
    i'
      -- Something like a division-rounding-up, but accomodating negative
      -- operands.
      let divRounding :: m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    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
      SubExp
    n <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "n" (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 (Exp (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    forall {m :: * -> *}.
    MonadBuilder m =>
    m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    divRounding (SubExp -> InternaliseM (Exp (Rep InternaliseM))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    toExp SubExp
    j_m_i) (SubExp -> InternaliseM (Exp (Rep InternaliseM))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    toExp SubExp
    s')

      SubExp
    zero_stride <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "zero_stride" (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
    CmpEq PrimType
    int64) SubExp
    s_sign SubExp
    zero
      SubExp
    nonzero_stride <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "nonzero_stride" (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
    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.
      SubExp
    empty_slice <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "empty_slice" (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
    CmpEq PrimType
    int64) SubExp
    n SubExp
    zero

      SubExp
    m <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "m" (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
    Sub IntType
    Int64 Overflow
    I.OverflowWrap) SubExp
    n SubExp
    one
      SubExp
    m_t_s <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "m_t_s" (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
    Mul IntType
    Int64 Overflow
    I.OverflowWrap) SubExp
    m SubExp
    s'
      SubExp
    i_p_m_t_s <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "i_p_m_t_s" (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
    Add IntType
    Int64 Overflow
    I.OverflowWrap) SubExp
    i' SubExp
    m_t_s
      SubExp
    zero_leq_i_p_m_t_s <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "zero_leq_i_p_m_t_s" (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 (IntType -> CmpOp
    I.CmpSle IntType
    Int64) SubExp
    zero SubExp
    i_p_m_t_s
      SubExp
    i_p_m_t_s_leq_w <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "i_p_m_t_s_leq_w" (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 (IntType -> CmpOp
    I.CmpSle IntType
    Int64) SubExp
    i_p_m_t_s SubExp
    w
      SubExp
    i_p_m_t_s_lth_w <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "i_p_m_t_s_leq_w" (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 (IntType -> CmpOp
    I.CmpSlt IntType
    Int64) SubExp
    i_p_m_t_s SubExp
    w

      SubExp
    zero_lte_i <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "zero_lte_i" (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 (IntType -> CmpOp
    I.CmpSle IntType
    Int64) SubExp
    zero SubExp
    i'
      SubExp
    i_lte_j <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "i_lte_j" (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 (IntType -> CmpOp
    I.CmpSle IntType
    Int64) SubExp
    i' SubExp
    j'
      SubExp
    forwards_ok <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "forwards_ok"
          (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAll [SubExp
    zero_lte_i, SubExp
    i_lte_j, SubExp
    zero_leq_i_p_m_t_s, SubExp
    i_p_m_t_s_lth_w]

      SubExp
    negone_lte_j <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "negone_lte_j" (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 (IntType -> CmpOp
    I.CmpSle IntType
    Int64) SubExp
    negone SubExp
    j'
      SubExp
    j_lte_i <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "j_lte_i" (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 (IntType -> CmpOp
    I.CmpSle IntType
    Int64) SubExp
    j' SubExp
    i'
      SubExp
    backwards_ok <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "backwards_ok"
          (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAll
            [SubExp
    negone_lte_j, SubExp
    j_lte_i, SubExp
    zero_leq_i_p_m_t_s, SubExp
    i_p_m_t_s_leq_w]

      SubExp
    slice_ok <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "slice_ok"
          (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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    backwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    backwards_ok])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    forwards_ok])

      SubExp
    ok_or_empty <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "ok_or_empty" (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.LogOr SubExp
    empty_slice SubExp
    slice_ok

      SubExp
    acceptable <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "slice_acceptable" (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
    nonzero_stride SubExp
    ok_or_empty

      let parts :: [ErrorMsgPart SubExp]
    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
    ":"]
      (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp -> SubExp -> SubExp -> DimIndex SubExp
    forall d. d -> d -> d -> DimIndex d
    I.DimSlice SubExp
    i' SubExp
    n SubExp
    s', SubExp
    acceptable, [ErrorMsgPart SubExp]
    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)

    internaliseScanOrReduce ::
      String ->
      String ->
      (SubExp -> I.Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) ->
      (E.Exp, E.Exp, E.Exp) ->
      InternaliseM [SubExp]
    internaliseScanOrReduce :: [Char]
    -> [Char]
    -> (SubExp
        -> Lambda SOACS
        -> [SubExp]
        -> [VName]
        -> InternaliseM (SOAC SOACS))
    -> (Exp, Exp, Exp)
    -> InternaliseM [SubExp]
    internaliseScanOrReduce [Char]
    desc [Char]
    what SubExp
    -> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
    f (Exp
    lam, Exp
    ne, Exp
    arr) = do
      [VName]
    arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars ([Char]
    what [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_arr") Exp
    arr
      [SubExp]
    nes <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp ([Char]
    what [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_ne") Exp
    ne
      [SubExp]
    nes' <- [(SubExp, VName)]
    -> ((SubExp, VName) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([SubExp] -> [VName] -> [(SubExp, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    nes [VName]
    arrs) (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
    -> ((SubExp, VName) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(SubExp
    ne', VName
    arr') -> do
        TypeBase Shape NoUniqueness
    rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. Int -> TypeBase Shape u -> TypeBase Shape u
    I.stripArray Int
    1 (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    arr'
        ErrorMsg SubExp
    -> TypeBase Shape NoUniqueness
    -> [Char]
    -> SubExp
    -> InternaliseM SubExp
    ensureShape
          ErrorMsg SubExp
    "Row shape of input array does not match shape of neutral element"
          TypeBase Shape NoUniqueness
    rowtype
          ([Char]
    what [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_ne_right_shape")
          SubExp
    ne'
      [TypeBase Shape NoUniqueness]
    nests <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType [SubExp]
    nes'
      [TypeBase Shape NoUniqueness]
    arrts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arrs
      Lambda SOACS
    lam' <- InternaliseLambda
    -> Exp
    -> [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM (Lambda SOACS)
    internaliseFoldLambda InternaliseLambda
    internaliseLambda Exp
    lam [TypeBase Shape NoUniqueness]
    nests [TypeBase Shape NoUniqueness]
    arrts
      SubExp
    w <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 ([TypeBase Shape NoUniqueness] -> SubExp)
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arrs
      [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    desc (Exp SOACS -> InternaliseM [SubExp])
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    I.Op (SOAC SOACS -> InternaliseM [SubExp])
    -> InternaliseM (SOAC SOACS) -> InternaliseM [SubExp]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< SubExp
    -> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
    f SubExp
    w Lambda SOACS
    lam' [SubExp]
    nes' [VName]
    arrs

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

      -- reshape neutral element to have same size as the destination array
      [SubExp]
    ne_shp <- [(SubExp, VName)]
    -> ((SubExp, VName) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([SubExp] -> [VName] -> [(SubExp, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    ne' [VName]
    hist') (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
    -> ((SubExp, VName) -> InternaliseM SubExp)
    -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(SubExp
    n, VName
    h) -> do
        TypeBase Shape NoUniqueness
    rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. Int -> TypeBase Shape u -> TypeBase Shape u
    I.stripArray Int
    1 (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    h
        ErrorMsg SubExp
    -> TypeBase Shape NoUniqueness
    -> [Char]
    -> SubExp
    -> InternaliseM SubExp
    ensureShape
          ErrorMsg SubExp
    "Row shape of destination array does not match shape of neutral element"
          TypeBase Shape NoUniqueness
    rowtype
          [Char]
    "hist_ne_right_shape"
          SubExp
    n
      [TypeBase Shape NoUniqueness]
    ne_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType [SubExp]
    ne_shp
      [TypeBase Shape NoUniqueness]
    his_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. Int -> TypeBase Shape u -> TypeBase Shape u
    I.stripArray (Int
    dim Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1)) (InternaliseM (TypeBase Shape NoUniqueness)
     -> InternaliseM (TypeBase Shape NoUniqueness))
    -> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> VName
    -> InternaliseM (TypeBase Shape NoUniqueness)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType) [VName]
    hist'
      Lambda SOACS
    op' <- InternaliseLambda
    -> Exp
    -> [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM (Lambda SOACS)
    internaliseFoldLambda InternaliseLambda
    internaliseLambda Exp
    op [TypeBase Shape NoUniqueness]
    ne_ts [TypeBase Shape NoUniqueness]
    his_ts

      -- reshape return type of bucket function to have same size as neutral element
      -- (modulo the index)
      [Param (TypeBase Shape NoUniqueness)]
    bucket_params <- Int
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM Int
    dim ([Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "bucket_p" (TypeBase Shape NoUniqueness
     -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall a b. (a -> b) -> a -> b
    $ PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64)
      [Param (TypeBase Shape NoUniqueness)]
    img_params <- (TypeBase Shape NoUniqueness
     -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    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 ([Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "img_p" (TypeBase Shape NoUniqueness
     -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
    -> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType) ([TypeBase Shape NoUniqueness]
     -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    img'
      let params :: [Param (TypeBase Shape NoUniqueness)]
    params = [Param (TypeBase Shape NoUniqueness)]
    bucket_params [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. [a] -> [a] -> [a]
    ++ [Param (TypeBase Shape NoUniqueness)]
    img_params
          rettype :: [TypeBase Shape NoUniqueness]
    rettype = Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
    forall a. Int -> a -> [a]
    replicate Int
    dim (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64) [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a. [a] -> [a] -> [a]
    ++ [TypeBase Shape NoUniqueness]
    ne_ts
          body :: Body SOACS
    body = Stms SOACS -> Result -> Body SOACS
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms SOACS
    forall a. Monoid a => a
    mempty (Result -> Body SOACS) -> Result -> Body SOACS
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Result
    varsRes ([VName] -> Result) -> [VName] -> Result
    forall a b. (a -> b) -> a -> b
    $ (Param (TypeBase Shape NoUniqueness) -> VName)
    -> [Param (TypeBase Shape NoUniqueness)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName [Param (TypeBase Shape NoUniqueness)]
    params
      Lambda SOACS
    lam' <-
        [LParam (Rep InternaliseM)]
    -> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [Param (TypeBase Shape NoUniqueness)]
    [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
    -> [TypeBase Shape NoUniqueness] -> Result -> InternaliseM Result
    ensureResultShape
            ErrorMsg SubExp
    "Row shape of value array does not match row shape of hist target"
            [TypeBase Shape NoUniqueness]
    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

      -- get sizes of histogram and image arrays
      Shape
    shape_hist <- [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape ([SubExp] -> Shape)
    -> (TypeBase Shape NoUniqueness -> [SubExp])
    -> TypeBase Shape NoUniqueness
    -> Shape
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take Int
    dim ([SubExp] -> [SubExp])
    -> (TypeBase Shape NoUniqueness -> [SubExp])
    -> TypeBase Shape NoUniqueness
    -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims (TypeBase Shape NoUniqueness -> Shape)
    -> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Shape
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType ([VName] -> VName
    forall a. HasCallStack => [a] -> a
    head [VName]
    hist')
      SubExp
    w_img <- Int -> TypeBase Shape NoUniqueness -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    I.arraySize Int
    0 (TypeBase Shape NoUniqueness -> SubExp)
    -> InternaliseM (TypeBase Shape NoUniqueness)
    -> InternaliseM SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType ([VName] -> VName
    forall a. HasCallStack => [a] -> a
    head [VName]
    img')

      [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    desc (Exp SOACS -> InternaliseM [SubExp])
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    I.Op (SOAC SOACS -> InternaliseM [SubExp])
    -> SOAC SOACS -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $
        SubExp -> [VName] -> [HistOp SOACS] -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
    I.Hist SubExp
    w_img ([VName]
    buckets' [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    img') [Shape
    -> SubExp -> [VName] -> [SubExp] -> Lambda SOACS -> HistOp SOACS
    forall rep.
    Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
    HistOp Shape
    shape_hist SubExp
    rf' [VName]
    hist' [SubExp]
    ne_shp Lambda SOACS
    op'] Lambda SOACS
    lam'

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

      VName
    acc_cert_v <- [Char] -> InternaliseM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "acc_cert"
      [TypeBase Shape NoUniqueness]
    dest_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    dest'
      let dest_w :: SubExp
    dest_w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    dest_ts
          acc_t :: TypeBase Shape NoUniqueness
    acc_t = VName
    -> Shape
    -> [TypeBase Shape NoUniqueness]
    -> NoUniqueness
    -> TypeBase Shape NoUniqueness
    forall shape u.
    VName
    -> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
    Acc VName
    acc_cert_v ([SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape [SubExp
    dest_w]) ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType [TypeBase Shape NoUniqueness]
    dest_ts) NoUniqueness
    NoUniqueness
      Param (TypeBase Shape NoUniqueness)
    acc_p <- [Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "acc_p" TypeBase Shape NoUniqueness
    acc_t
      Lambda SOACS
    withacc_lam <- [LParam (Rep InternaliseM)]
    -> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [Attrs
    -> VName
    -> TypeBase Shape NoUniqueness
    -> Param (TypeBase Shape NoUniqueness)
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    forall a. Monoid a => a
    mempty VName
    acc_cert_v (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    I.Unit), Param (TypeBase Shape NoUniqueness)
    LParam (Rep InternaliseM)
    acc_p] (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
    -> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $ do
        [TypeBase Shape NoUniqueness]
    bs_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    bs'
        Lambda SOACS
    lam' <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
    internaliseLambdaCoerce Exp
    lam ([TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS))
    -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
    forall a b. (a -> b) -> a -> b
    $ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType ([TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness])
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    paramType Param (TypeBase Shape NoUniqueness)
    acc_p TypeBase Shape NoUniqueness
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a. a -> [a] -> [a]
    : [TypeBase Shape NoUniqueness]
    bs_ts
        let w :: SubExp
    w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    bs_ts
        ([SubExp] -> Result)
    -> InternaliseM [SubExp] -> InternaliseM Result
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [SubExp] -> Result
    subExpsRes (InternaliseM [SubExp] -> InternaliseM Result)
    -> (Exp SOACS -> InternaliseM [SubExp])
    -> Exp SOACS
    -> InternaliseM Result
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    "acc_res" (Exp SOACS -> InternaliseM Result)
    -> Exp SOACS -> InternaliseM Result
    forall a b. (a -> b) -> a -> b
    $
          Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    I.Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $
            SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    I.Screma SubExp
    w (Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName Param (TypeBase Shape NoUniqueness)
    acc_p VName -> [VName] -> [VName]
    forall a. a -> [a] -> [a]
    : [VName]
    bs') (Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    I.mapSOAC Lambda SOACS
    lam')

      Maybe (Lambda SOACS, [SubExp])
    op' <-
        case Maybe (Exp, Exp)
    op of
          Just (Exp
    op_lam, Exp
    ne) -> do
            [SubExp]
    ne' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "hist_ne" Exp
    ne
            [TypeBase Shape NoUniqueness]
    ne_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType [SubExp]
    ne'
            ([Param (TypeBase Shape NoUniqueness)]
    lam_params, Body SOACS
    lam_body, [TypeBase Shape NoUniqueness]
    lam_rettype) <-
              InternaliseLambda
    internaliseLambda Exp
    op_lam ([TypeBase Shape NoUniqueness]
     -> InternaliseM
          ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
    forall a b. (a -> b) -> a -> b
    $ [TypeBase Shape NoUniqueness]
    ne_ts [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a. [a] -> [a] -> [a]
    ++ [TypeBase Shape NoUniqueness]
    ne_ts
            Param (TypeBase Shape NoUniqueness)
    idxp <- [Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "idx" (TypeBase Shape NoUniqueness
     -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall a b. (a -> b) -> a -> b
    $ PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64
            let op_lam' :: Lambda SOACS
    op_lam' = [LParam SOACS]
    -> [TypeBase Shape NoUniqueness] -> Body SOACS -> Lambda SOACS
    forall rep.
    [LParam rep]
    -> [TypeBase Shape NoUniqueness] -> Body rep -> Lambda rep
    I.Lambda (Param (TypeBase Shape NoUniqueness)
    idxp Param (TypeBase Shape NoUniqueness)
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. a -> [a] -> [a]
    : [Param (TypeBase Shape NoUniqueness)]
    lam_params) [TypeBase Shape NoUniqueness]
    lam_rettype Body SOACS
    lam_body
            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])
     -> InternaliseM (Maybe (Lambda SOACS, [SubExp])))
    -> Maybe (Lambda SOACS, [SubExp])
    -> InternaliseM (Maybe (Lambda SOACS, [SubExp]))
    forall a b. (a -> b) -> a -> b
    $ (Lambda SOACS, [SubExp]) -> Maybe (Lambda SOACS, [SubExp])
    forall a. a -> Maybe a
    Just (Lambda SOACS
    op_lam', [SubExp]
    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

      SubExp
    destw <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 ([TypeBase Shape NoUniqueness] -> SubExp)
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    dest'
      ([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) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var) (InternaliseM [VName] -> InternaliseM [SubExp])
    -> InternaliseM [VName] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    desc (Exp (Rep InternaliseM) -> InternaliseM [VName])
    -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $
          [WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
    forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
    WithAcc [([SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape [SubExp
    destw], [VName]
    dest', Maybe (Lambda SOACS, [SubExp])
    op')] Lambda SOACS
    withacc_lam

    internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp
    internaliseExp1 :: [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    desc Exp
    e = do
      [SubExp]
    vs <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
      case [SubExp]
    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 :: String -> E.Exp -> InternaliseM (I.SubExp, IntType)
    internaliseSizeExp :: [Char] -> Exp -> InternaliseM (SubExp, IntType)
    internaliseSizeExp [Char]
    s Exp
    e = do
      SubExp
    e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    s Exp
    e
      case Exp -> StructType
    E.typeOf Exp
    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"

    internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
    internaliseExpToVars :: [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    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 SubExp -> InternaliseM VName
    asIdent ([SubExp] -> InternaliseM [VName])
    -> InternaliseM [SubExp] -> InternaliseM [VName]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
      where
        asIdent :: SubExp -> InternaliseM VName
    asIdent (I.Var VName
    v) = VName -> InternaliseM VName
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    v
        asIdent SubExp
    se = [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    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

    internaliseOperation ::
      String ->
      E.Exp ->
      (I.VName -> InternaliseM I.BasicOp) ->
      InternaliseM [I.SubExp]
    internaliseOperation :: [Char]
    -> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
    internaliseOperation [Char]
    s Exp
    e VName -> InternaliseM BasicOp
    op = do
      [VName]
    vs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    s Exp
    e
      (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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    s (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)
    -> (VName -> InternaliseM BasicOp) -> VName -> InternaliseM SubExp
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< VName -> InternaliseM BasicOp
    op) [VName]
    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
      SubExp
    zero <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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)
      SubExp
    nonzero <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "nonzero" (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
    UnOp (PrimType -> UnOp
    I.Neg PrimType
    I.Bool) SubExp
    zero
      Certs
    c <- [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "nonzero_cert" SubExp
    nonzero ErrorMsg SubExp
    "division by zero"
      Certs -> InternaliseM a -> InternaliseM a
    forall a. Certs -> InternaliseM a -> InternaliseM a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying Certs
    c InternaliseM a
    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
      SubExp
    nonnegative <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
      Certs
    c <- [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "nonzero_cert" SubExp
    nonnegative ErrorMsg SubExp
    "negative exponent"
      Certs -> InternaliseM a -> InternaliseM a
    forall a. Certs -> InternaliseM a -> InternaliseM a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying Certs
    c InternaliseM a
    m

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

    -- Relational operators for booleans.
    internaliseBinOp [Char]
    desc BinOp
    E.Less SubExp
    x SubExp
    y PrimType
    E.Bool PrimType
    _ =
      [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
    simpleCmpOp [Char]
    desc CmpOp
    I.CmpLlt SubExp
    x SubExp
    y
    internaliseBinOp [Char]
    desc BinOp
    E.Leq SubExp
    x SubExp
    y PrimType
    E.Bool PrimType
    _ =
      [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
    simpleCmpOp [Char]
    desc CmpOp
    I.CmpLle SubExp
    x SubExp
    y
    internaliseBinOp [Char]
    desc BinOp
    E.Greater SubExp
    x SubExp
    y PrimType
    E.Bool PrimType
    _ =
      [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
    simpleCmpOp [Char]
    desc CmpOp
    I.CmpLlt SubExp
    y SubExp
    x -- Note the swapped x and y
    internaliseBinOp [Char]
    desc BinOp
    E.Geq SubExp
    x SubExp
    y PrimType
    E.Bool PrimType
    _ =
      [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
    simpleCmpOp [Char]
    desc CmpOp
    I.CmpLle SubExp
    y SubExp
    x -- Note the swapped x and y
    internaliseBinOp [Char]
    _ 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 ::
      String ->
      I.BinOp ->
      I.SubExp ->
      I.SubExp ->
      InternaliseM [I.SubExp]
    simpleBinOp :: [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
    simpleBinOp [Char]
    desc BinOp
    bop SubExp
    x SubExp
    y =
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 ::
      String ->
      I.CmpOp ->
      I.SubExp ->
      I.SubExp ->
      InternaliseM [I.SubExp]
    simpleCmpOp :: [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
    simpleCmpOp [Char]
    desc CmpOp
    op SubExp
    x SubExp
    y =
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 [ExtType]
    bodyExtType (Body BodyDec SOACS
    _ Stms SOACS
    stms Result
    res) =
      [VName] -> [ExtType] -> [ExtType]
    existentialiseExtTypes (Scope SOACS -> [VName]
    forall k a. Map k a -> [k]
    M.keys Scope SOACS
    stmsscope) ([ExtType] -> [ExtType])
    -> ([TypeBase Shape NoUniqueness] -> [ExtType])
    -> [TypeBase Shape NoUniqueness]
    -> [ExtType]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [TypeBase Shape NoUniqueness] -> [ExtType]
    forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
    staticShapes
        ([TypeBase Shape NoUniqueness] -> [ExtType])
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM [ExtType]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ExtendedScope SOACS InternaliseM [TypeBase Shape NoUniqueness]
    -> Scope SOACS -> InternaliseM [TypeBase Shape NoUniqueness]
    forall rep (m :: * -> *) a.
    ExtendedScope rep m a -> Scope rep -> m a
    extendedScope ((SubExpRes
     -> ExtendedScope SOACS InternaliseM (TypeBase Shape NoUniqueness))
    -> Result
    -> ExtendedScope SOACS InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExpRes -> m (TypeBase Shape NoUniqueness)
    subExpResType Result
    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 :: InternaliseLambda
    internaliseLambda :: InternaliseLambda
    internaliseLambda (E.Parens Exp
    e SrcLoc
    _) [TypeBase Shape NoUniqueness]
    rowtypes =
      InternaliseLambda
    internaliseLambda Exp
    e [TypeBase Shape NoUniqueness]
    rowtypes
    internaliseLambda (E.Lambda [PatBase Info VName ParamType]
    params Exp
    body Maybe (TypeExp Exp VName)
    _ (Info (RetType [VName]
    _ TypeBase Exp Uniqueness
    rettype)) SrcLoc
    _) [TypeBase Shape NoUniqueness]
    rowtypes =
      [PatBase Info VName ParamType]
    -> [TypeBase Shape NoUniqueness]
    -> ([LParam SOACS]
        -> InternaliseM
             ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
    forall a.
    [PatBase Info VName ParamType]
    -> [TypeBase Shape NoUniqueness]
    -> ([LParam SOACS] -> InternaliseM a)
    -> InternaliseM a
    bindingLambdaParams [PatBase Info VName ParamType]
    params [TypeBase Shape NoUniqueness]
    rowtypes (([LParam SOACS]
      -> InternaliseM
           ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
     -> InternaliseM
          ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
    -> ([LParam SOACS]
        -> InternaliseM
             ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
    forall a b. (a -> b) -> a -> b
    $ \[LParam SOACS]
    params' -> do
        Body SOACS
    body' <- [Char] -> Exp -> InternaliseM (Body SOACS)
    internaliseBody [Char]
    "lam" Exp
    body
        [TypeBase Shape NoUniqueness]
    rettype' <- TypeBase Exp Uniqueness
    -> [ExtType] -> InternaliseM [TypeBase Shape NoUniqueness]
    forall shape u.
    TypeBase Exp Uniqueness
    -> [TypeBase shape u] -> InternaliseM [TypeBase Shape NoUniqueness]
    internaliseLambdaReturnType TypeBase Exp Uniqueness
    rettype ([ExtType] -> InternaliseM [TypeBase Shape NoUniqueness])
    -> InternaliseM [ExtType]
    -> InternaliseM [TypeBase Shape NoUniqueness]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Body SOACS -> InternaliseM [ExtType]
    bodyExtType Body SOACS
    body'
        ([Param (TypeBase Shape NoUniqueness)], Body SOACS,
     [TypeBase Shape NoUniqueness])
    -> InternaliseM
         ([Param (TypeBase Shape NoUniqueness)], Body SOACS,
          [TypeBase Shape NoUniqueness])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Param (TypeBase Shape NoUniqueness)]
    [LParam SOACS]
    params', Body SOACS
    body', [TypeBase Shape NoUniqueness]
    rettype')
    internaliseLambda Exp
    e [TypeBase Shape NoUniqueness]
    _ = [Char]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
    forall a. HasCallStack => [Char] -> a
    error ([Char]
     -> InternaliseM
          ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
    -> [Char]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
    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 NoUniqueness] -> InternaliseM (Lambda SOACS)
    internaliseLambdaCoerce Exp
    lam [TypeBase Shape NoUniqueness]
    argtypes = do
      ([Param (TypeBase Shape NoUniqueness)]
    params, Body SOACS
    body, [TypeBase Shape NoUniqueness]
    rettype) <- InternaliseLambda
    internaliseLambda Exp
    lam [TypeBase Shape NoUniqueness]
    argtypes
      [LParam (Rep InternaliseM)]
    -> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [Param (TypeBase Shape NoUniqueness)]
    [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
    -> [TypeBase Shape NoUniqueness] -> 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
    "unexpected lambda result size"])
          [TypeBase Shape NoUniqueness]
    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

    -- | Overloaded operators are treated here.
    isOverloadedFunction ::
      E.QualName VName ->
      String ->
      Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
    isOverloadedFunction :: QualName VName
    -> [Char]
    -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
    isOverloadedFunction QualName VName
    qname [Char]
    desc = do
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qname) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag
      [Char] -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
    handle ([Char]
     -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
    -> [Char]
    -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ VName -> [Char]
    baseString (VName -> [Char]) -> VName -> [Char]
    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 :: [Char] -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
    handle [Char]
    op
          | Just SubExp -> InternaliseM [SubExp]
    cmp_f <- [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
    isEqlOp [Char]
    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
              [SubExp]
    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'
              SubExp -> InternaliseM [SubExp]
    cmp_f (SubExp -> InternaliseM [SubExp])
    -> InternaliseM SubExp -> InternaliseM [SubExp]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "eq" (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAll [SubExp]
    rs
          where
            isEqlOp :: [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
    isEqlOp [Char]
    "!=" = (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 ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 [Char]
    "==" = (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 [Char]
    _ = Maybe (SubExp -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

            doComparison :: SubExp -> SubExp -> InternaliseM SubExp
    doComparison SubExp
    x SubExp
    y = do
              TypeBase Shape NoUniqueness
    x_t <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType SubExp
    x
              TypeBase Shape NoUniqueness
    y_t <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType SubExp
    y
              case TypeBase Shape NoUniqueness
    x_t of
                I.Prim PrimType
    t -> [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 NoUniqueness
    _ -> do
                  let x_dims :: [SubExp]
    x_dims = TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims TypeBase Shape NoUniqueness
    x_t
                      y_dims :: [SubExp]
    y_dims = TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims TypeBase Shape NoUniqueness
    y_t
                  [SubExp]
    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) ->
                    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
                  SubExp
    shapes_match <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "shapes_match" (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAll [SubExp]
    dims_match
                  let compare_elems_body :: InternaliseM (Body SOACS)
    compare_elems_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
                        -- Flatten both x and y.
                        SubExp
    x_num_elems <-
                          [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
                        VName
    x' <- [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "x" (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) VName)
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    I.SubExp SubExp
    x
                        VName
    y' <- [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "x" (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) VName)
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    I.SubExp SubExp
    y
                        VName
    x_flat <-
                          [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "x_flat" (Exp SOACS -> BuilderT SOACS (State VNameSource) VName)
    -> (BasicOp -> Exp SOACS)
    -> BasicOp
    -> BuilderT SOACS (State VNameSource) VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp SOACS
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> BuilderT SOACS (State VNameSource) VName)
    -> BasicOp -> BuilderT SOACS (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $
                            VName -> NewShape SubExp -> BasicOp
    I.Reshape VName
    x' (Shape -> Shape -> NewShape SubExp
    forall old new.
    ArrayShape old =>
    old -> ShapeBase new -> NewShape new
    reshapeAll (TypeBase Shape NoUniqueness -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    I.arrayShape TypeBase Shape NoUniqueness
    x_t) ([SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape [SubExp
    x_num_elems]))
                        VName
    y_flat <-
                          [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "y_flat" (Exp SOACS -> BuilderT SOACS (State VNameSource) VName)
    -> (BasicOp -> Exp SOACS)
    -> BasicOp
    -> BuilderT SOACS (State VNameSource) VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp SOACS
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> BuilderT SOACS (State VNameSource) VName)
    -> BasicOp -> BuilderT SOACS (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $
                            VName -> NewShape SubExp -> BasicOp
    I.Reshape VName
    y' (Shape -> Shape -> NewShape SubExp
    forall old new.
    ArrayShape old =>
    old -> ShapeBase new -> NewShape new
    reshapeAll (TypeBase Shape NoUniqueness -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    I.arrayShape TypeBase Shape NoUniqueness
    x_t) ([SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape [SubExp
    x_num_elems]))

                        -- Compare the elements.
                        Lambda SOACS
    cmp_lam <- CmpOp
    -> BuilderT
         SOACS
         (State VNameSource)
         (Lambda (Rep (BuilderT SOACS (State VNameSource))))
    forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    CmpOp -> m (Lambda (Rep m))
    cmpOpLambda (CmpOp
     -> BuilderT
          SOACS
          (State VNameSource)
          (Lambda (Rep (BuilderT SOACS (State VNameSource)))))
    -> CmpOp
    -> BuilderT
         SOACS
         (State VNameSource)
         (Lambda (Rep (BuilderT SOACS (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ PrimType -> CmpOp
    I.CmpEq (TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    x_t)
                        VName
    cmps <-
                          [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "cmps" (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) VName)
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $
                            Op (Rep (BuilderT SOACS (State VNameSource)))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. Op rep -> Exp rep
    I.Op (Op (Rep (BuilderT SOACS (State VNameSource)))
     -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> Op (Rep (BuilderT SOACS (State VNameSource)))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $
                              SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    I.Screma SubExp
    x_num_elems [VName
    x_flat, VName
    y_flat] (Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    I.mapSOAC Lambda SOACS
    cmp_lam)

                        -- Check that all were equal.
                        Lambda SOACS
    and_lam <- BinOp
    -> PrimType
    -> BuilderT
         SOACS
         (State VNameSource)
         (Lambda (Rep (BuilderT SOACS (State VNameSource))))
    forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    BinOp -> PrimType -> m (Lambda (Rep m))
    binOpLambda BinOp
    I.LogAnd PrimType
    I.Bool
                        ScremaForm SOACS
    reduce <- [Reduce SOACS]
    -> BuilderT SOACS (State VNameSource) (ScremaForm SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Reduce rep] -> m (ScremaForm rep)
    I.reduceSOAC [Commutativity -> Lambda SOACS -> [SubExp] -> Reduce SOACS
    forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    Reduce Commutativity
    Commutative Lambda SOACS
    and_lam [Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    True]]
                        SubExp
    all_equal <- [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "all_equal" (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall a b. (a -> b) -> a -> b
    $ Op (Rep (BuilderT SOACS (State VNameSource)))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. Op rep -> Exp rep
    I.Op (Op (Rep (BuilderT SOACS (State VNameSource)))
     -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> Op (Rep (BuilderT SOACS (State VNameSource)))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    I.Screma SubExp
    x_num_elems [VName
    cmps] ScremaForm SOACS
    reduce
                        Result -> Builder SOACS Result
    forall a. a -> BuilderT SOACS (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Result -> Builder SOACS Result) -> Result -> Builder SOACS Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Result
    subExpsRes [SubExp
    all_equal]

                  [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "arrays_equal"
                    (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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    shapes_match) InternaliseM (Body (Rep InternaliseM))
    InternaliseM (Body SOACS)
    compare_elems_body ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    False])
        handle [Char]
    name
          | Just BinOp
    bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (([Char]
    name ==) ([Char] -> Bool) -> (BinOp -> [Char]) -> BinOp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BinOp -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString) [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)) ->
                    [Char]
    -> BinOp
    -> SubExp
    -> SubExp
    -> PrimType
    -> PrimType
    -> InternaliseM [SubExp]
    internaliseBinOp [Char]
    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 [Char]
    _ = Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

    -- | 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 (String -> InternaliseM [SubExp])
    isIntrinsicFunction :: QualName VName -> [Exp] -> Maybe ([Char] -> 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 -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qname) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag
      let handlers :: [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
    handlers =
            [ [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall {a}.
    (Eq a, IsString a) =>
    [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleSign,
              [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall {f :: * -> *}.
    Applicative f =>
    [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
    handleOps,
              [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall {a}.
    (IsString a, Eq a) =>
    [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleSOACs,
              [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall {a}.
    (Eq a, IsString a) =>
    [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleAccs,
              [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall {a}.
    (IsString a, Eq a) =>
    [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleAD,
              [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall {a}.
    (IsString a, Eq a) =>
    [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleRest
            ]
      [Maybe ([Char] -> InternaliseM [SubExp])]
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall (t :: * -> *) (m :: * -> *) a.
    (Foldable t, MonadPlus m) =>
    t (m a) -> m a
    msum [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    h [Exp]
    args ([Char] -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ VName -> [Char]
    baseString (VName -> [Char]) -> VName -> [Char]
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    qname | [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
    h <- [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
    handlers]
      where
        handleSign :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleSign [Exp
    x] a
    "sign_i8" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toSigned IntType
    I.Int8 Exp
    x
        handleSign [Exp
    x] a
    "sign_i16" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toSigned IntType
    I.Int16 Exp
    x
        handleSign [Exp
    x] a
    "sign_i32" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toSigned IntType
    I.Int32 Exp
    x
        handleSign [Exp
    x] a
    "sign_i64" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toSigned IntType
    I.Int64 Exp
    x
        handleSign [Exp
    x] a
    "unsign_i8" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toUnsigned IntType
    I.Int8 Exp
    x
        handleSign [Exp
    x] a
    "unsign_i16" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toUnsigned IntType
    I.Int16 Exp
    x
        handleSign [Exp
    x] a
    "unsign_i32" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toUnsigned IntType
    I.Int32 Exp
    x
        handleSign [Exp
    x] a
    "unsign_i64" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toUnsigned IntType
    I.Int64 Exp
    x
        handleSign [Exp]
    _ a
    _ = Maybe ([Char] -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

        handleOps :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
    handleOps [Exp
    x] [Char]
    s
          | Just UnOp
    unop <- (UnOp -> Bool) -> [UnOp] -> Maybe UnOp
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (([Char] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Char]
    s) ([Char] -> Bool) -> (UnOp -> [Char]) -> UnOp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . UnOp -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString) [UnOp]
    allUnOps = ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM (f SubExp))
     -> Maybe ([Char] -> InternaliseM (f SubExp)))
    -> ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
              SubExp
    x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              (SubExp -> f SubExp)
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> f SubExp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (InternaliseM SubExp -> InternaliseM (f SubExp))
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> a -> b
    $ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 UnOp
    unop SubExp
    x'
        handleOps [TupLit [Exp
    x, Exp
    y] SrcLoc
    _] [Char]
    s
          | Just BinOp
    bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (([Char] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Char]
    s) ([Char] -> Bool) -> (BinOp -> [Char]) -> BinOp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BinOp -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString) [BinOp]
    allBinOps = ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM (f SubExp))
     -> Maybe ([Char] -> InternaliseM (f SubExp)))
    -> ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
              SubExp
    x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              SubExp
    y' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "y" Exp
    y
              (SubExp -> f SubExp)
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> f SubExp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (InternaliseM SubExp -> InternaliseM (f SubExp))
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> a -> b
    $ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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'
          | Just CmpOp
    cmp <- (CmpOp -> Bool) -> [CmpOp] -> Maybe CmpOp
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (([Char] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Char]
    s) ([Char] -> Bool) -> (CmpOp -> [Char]) -> CmpOp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . CmpOp -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString) [CmpOp]
    allCmpOps = ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM (f SubExp))
     -> Maybe ([Char] -> InternaliseM (f SubExp)))
    -> ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
              SubExp
    x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              SubExp
    y' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "y" Exp
    y
              (SubExp -> f SubExp)
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> f SubExp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (InternaliseM SubExp -> InternaliseM (f SubExp))
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> a -> b
    $ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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
    cmp SubExp
    x' SubExp
    y'
        handleOps [Exp
    x] [Char]
    s
          | Just ConvOp
    conv <- (ConvOp -> Bool) -> [ConvOp] -> Maybe ConvOp
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (([Char] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Char]
    s) ([Char] -> Bool) -> (ConvOp -> [Char]) -> ConvOp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ConvOp -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString) [ConvOp]
    allConvOps = ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM (f SubExp))
     -> Maybe ([Char] -> InternaliseM (f SubExp)))
    -> ([Char] -> InternaliseM (f SubExp))
    -> Maybe ([Char] -> InternaliseM (f SubExp))
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
              SubExp
    x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              (SubExp -> f SubExp)
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> f SubExp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (InternaliseM SubExp -> InternaliseM (f SubExp))
    -> InternaliseM SubExp -> InternaliseM (f SubExp)
    forall a b. (a -> b) -> a -> b
    $ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 ConvOp
    conv SubExp
    x'
        handleOps [Exp]
    _ [Char]
    _ = Maybe ([Char] -> InternaliseM (f SubExp))
    forall a. Maybe a
    Nothing

        handleSOACs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleSOACs [Exp
    lam, Exp
    arr] a
    "map" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [VName]
    arr' <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "map_arr" Exp
    arr
          [TypeBase Shape NoUniqueness]
    arr_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arr'
          Lambda SOACS
    lam' <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
    internaliseLambdaCoerce Exp
    lam ([TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS))
    -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
    forall a b. (a -> b) -> a -> b
    $ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType [TypeBase Shape NoUniqueness]
    arr_ts
          let w :: SubExp
    w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    arr_ts
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
    -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall rep. Op rep -> Exp rep
    I.Op (Op (Rep InternaliseM) -> Exp (Rep InternaliseM))
    -> Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    I.Screma SubExp
    w [VName]
    arr' (Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    I.mapSOAC Lambda SOACS
    lam')
        handleSOACs [Exp
    k, Exp
    lam, Exp
    arr] a
    "partition" = do
          Int
    k' <- Int32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (Int32 -> Int) -> Maybe Int32 -> Maybe Int
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> Maybe Int32
    forall {vn}. ExpBase Info vn -> Maybe Int32
    fromInt32 Exp
    k
          ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    _desc -> do
            [VName]
    arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "partition_input" Exp
    arr
            Lambda SOACS
    lam' <- InternaliseLambda
    -> Int -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
    internalisePartitionLambda InternaliseLambda
    internaliseLambda Int
    k' Exp
    lam ([SubExp] -> InternaliseM (Lambda SOACS))
    -> [SubExp] -> InternaliseM (Lambda SOACS)
    forall a b. (a -> b) -> a -> b
    $ (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var [VName]
    arrs
            ([SubExp] -> [SubExp] -> [SubExp])
    -> ([SubExp], [SubExp]) -> [SubExp]
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    (++) (([SubExp], [SubExp]) -> [SubExp])
    -> InternaliseM ([SubExp], [SubExp]) -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
    partitionWithSOACS (Int -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    k') Lambda SOACS
    lam' [VName]
    arrs
          where
            fromInt32 :: ExpBase Info vn -> Maybe Int32
    fromInt32 (Literal (SignedValue (Int32Value Int32
    k')) SrcLoc
    _) = Int32 -> Maybe Int32
    forall a. a -> Maybe a
    Just Int32
    k'
            fromInt32 (IntLit Integer
    k' (Info (E.Scalar (E.Prim (E.Signed IntType
    Int32)))) SrcLoc
    _) = Int32 -> Maybe Int32
    forall a. a -> Maybe a
    Just (Int32 -> Maybe Int32) -> Int32 -> Maybe Int32
    forall a b. (a -> b) -> a -> b
    $ Integer -> Int32
    forall a. Num a => Integer -> a
    fromInteger Integer
    k'
            fromInt32 ExpBase Info vn
    _ = Maybe Int32
    forall a. Maybe a
    Nothing
        handleSOACs [Exp
    lam, Exp
    ne, Exp
    arr] a
    "reduce" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          [Char]
    -> [Char]
    -> (SubExp
        -> Lambda SOACS
        -> [SubExp]
        -> [VName]
        -> InternaliseM (SOAC SOACS))
    -> (Exp, Exp, Exp)
    -> InternaliseM [SubExp]
    internaliseScanOrReduce [Char]
    desc [Char]
    "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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          [Char]
    -> [Char]
    -> (SubExp
        -> Lambda SOACS
        -> [SubExp]
        -> [VName]
        -> InternaliseM (SOAC SOACS))
    -> (Exp, Exp, Exp)
    -> InternaliseM [SubExp]
    internaliseScanOrReduce [Char]
    desc [Char]
    "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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          [Char]
    -> [Char]
    -> (SubExp
        -> Lambda SOACS
        -> [SubExp]
        -> [VName]
        -> InternaliseM (SOAC SOACS))
    -> (Exp, Exp, Exp)
    -> InternaliseM [SubExp]
    internaliseScanOrReduce [Char]
    desc [Char]
    "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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          Int
    -> [Char]
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> InternaliseM [SubExp]
    internaliseHist Int
    1 [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          Int
    -> [Char]
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> InternaliseM [SubExp]
    internaliseHist Int
    2 [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          Int
    -> [Char]
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> Exp
    -> InternaliseM [SubExp]
    internaliseHist Int
    3 [Char]
    desc Exp
    rf Exp
    dest Exp
    op Exp
    ne Exp
    buckets Exp
    img
        handleSOACs [Exp]
    _ a
    _ = Maybe ([Char] -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

        handleAccs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
    handleAccs [Exp
    dest, Exp
    f, Exp
    bs] a
    "scatter_stream" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          [Char]
    -> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
    internaliseStreamAcc [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          [Char]
    -> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
    internaliseStreamAcc [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          VName
    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
    <$> [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "acc" Exp
    acc
          SubExp
    i' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "acc_i" Exp
    i
          [SubExp]
    vs <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "acc_v" Exp
    v
          (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])
    -> InternaliseM SubExp -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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
    BasicOp (BasicOp -> Exp (Rep InternaliseM))
    -> BasicOp -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
    UpdateAcc Safety
    Safe VName
    acc' [SubExp
    i'] [SubExp]
    vs
        handleAccs [Exp]
    _ a
    _ = Maybe ([Char] -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

        handleAD :: [Exp] -> a -> Maybe ([Char] -> 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"] = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
              [SubExp]
    x' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "ad_x" Exp
    x
              [SubExp]
    v' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "ad_v" Exp
    v
              Lambda SOACS
    lam <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
    internaliseLambdaCoerce Exp
    f ([TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS))
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM (Lambda SOACS)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    x'
              ([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) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var) (InternaliseM [VName] -> InternaliseM [SubExp])
    -> (SOAC SOACS -> InternaliseM [VName])
    -> SOAC SOACS
    -> InternaliseM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    desc (Exp SOACS -> InternaliseM [VName])
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (SOAC SOACS -> InternaliseM [SubExp])
    -> SOAC SOACS -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $
                case a
    fname of
                  a
    "jvp2" -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    JVP [SubExp]
    x' [SubExp]
    v' Lambda SOACS
    lam
                  a
    _ -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    VJP [SubExp]
    x' [SubExp]
    v' Lambda SOACS
    lam
        handleAD [Exp]
    _ a
    _ = Maybe ([Char] -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

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

          let conc :: VName -> VName -> Exp SOACS
    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
          (Exp SOACS -> InternaliseM SubExp)
    -> [Exp SOACS] -> 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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    desc) ([Exp SOACS] -> InternaliseM [SubExp])
    -> [Exp SOACS] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ (VName -> VName -> Exp SOACS) -> [VName] -> [VName] -> [Exp SOACS]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith VName -> VName -> Exp SOACS
    conc [VName]
    xs [VName]
    ys
        handleRest [Exp
    e] a
    "transpose" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc ->
          [Char]
    -> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
    internaliseOperation [Char]
    desc Exp
    e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
    -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    v -> do
            Int
    r <- TypeBase Shape NoUniqueness -> Int
    forall shape u. ArrayShape shape => TypeBase shape u -> Int
    I.arrayRank (TypeBase Shape NoUniqueness -> Int)
    -> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Int
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    v
            BasicOp -> InternaliseM BasicOp
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (BasicOp -> InternaliseM BasicOp)
    -> BasicOp -> InternaliseM BasicOp
    forall a b. (a -> b) -> a -> b
    $ VName -> [Int] -> BasicOp
    I.Rearrange VName
    v ([Int
    1, Int
    0] [Int] -> [Int] -> [Int]
    forall a. [a] -> [a] -> [a]
    ++ [Int
    2 .. Int
    r Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1])
        handleRest [Exp
    x, Exp
    y] a
    "zip" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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
    <$> [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_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
    <*> [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_zip_y") Exp
    y
                )
        handleRest [Exp
    x] a
    "unzip" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    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 ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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
    =<< [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    desc Exp
    x
        handleRest [Exp
    arr, Exp
    offset, Exp
    n1, Exp
    s1, Exp
    n2, Exp
    s2] a
    "flat_index_2d" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [Char] -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
    flatIndexHelper [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [Char] -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
    flatUpdateHelper [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [Char] -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
    flatIndexHelper [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [Char] -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
    flatUpdateHelper [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [Char] -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
    flatIndexHelper [Char]
    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" = ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a. a -> Maybe a
    Just (([Char] -> InternaliseM [SubExp])
     -> Maybe ([Char] -> InternaliseM [SubExp]))
    -> ([Char] -> InternaliseM [SubExp])
    -> Maybe ([Char] -> InternaliseM [SubExp])
    forall a b. (a -> b) -> a -> b
    $ \[Char]
    desc -> do
          [Char] -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
    flatUpdateHelper [Char]
    desc Exp
    arr1 Exp
    offset [Exp
    s1, Exp
    s2, Exp
    s3, Exp
    s4] Exp
    arr2
        handleRest [Exp]
    _ a
    _ = Maybe ([Char] -> InternaliseM [SubExp])
    forall a. Maybe a
    Nothing

        toSigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
    toSigned IntType
    int_to Exp
    e [Char]
    desc = do
          SubExp
    e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "trunc_arg" Exp
    e
          case Exp -> StructType
    E.typeOf Exp
    e of
            E.Scalar (E.Prim PrimType
    E.Bool) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    e')
                  ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [IntType -> Integer -> SubExp
    intConst IntType
    int_to Integer
    1])
                  ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    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)) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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)) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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)) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 -> [Char] -> InternaliseM [SubExp]
    toUnsigned IntType
    int_to Exp
    e [Char]
    desc = do
          SubExp
    e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "trunc_arg" Exp
    e
          case Exp -> StructType
    E.typeOf Exp
    e of
            E.Scalar (E.Prim PrimType
    E.Bool) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    e')
                  ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [IntType -> Integer -> SubExp
    intConst IntType
    int_to Integer
    1])
                  ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    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)) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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)) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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)) ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    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"

        scatterF :: Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
    scatterF Int
    dim Exp
    a Exp
    si Exp
    v [Char]
    desc = do
          [VName]
    si' <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "write_arg_i" Exp
    si
          [VName]
    svs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "write_arg_v" Exp
    v
          [VName]
    sas <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "write_arg_a" Exp
    a

          SubExp
    si_w <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    I.arraysSize Int
    0 ([TypeBase Shape NoUniqueness] -> SubExp)
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    si'
          [TypeBase Shape NoUniqueness]
    sv_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    svs

          [VName]
    svs' <- [(VName, TypeBase Shape NoUniqueness)]
    -> ((VName, TypeBase Shape NoUniqueness) -> InternaliseM VName)
    -> InternaliseM [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([VName]
    -> [TypeBase Shape NoUniqueness]
    -> [(VName, TypeBase Shape NoUniqueness)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    svs [TypeBase Shape NoUniqueness]
    sv_ts) (((VName, TypeBase Shape NoUniqueness) -> InternaliseM VName)
     -> InternaliseM [VName])
    -> ((VName, TypeBase Shape NoUniqueness) -> InternaliseM VName)
    -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    sv, TypeBase Shape NoUniqueness
    sv_t) -> do
            let sv_shape :: Shape
    sv_shape = TypeBase Shape NoUniqueness -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    I.arrayShape TypeBase Shape NoUniqueness
    sv_t
                sv_w :: SubExp
    sv_w = Int -> TypeBase Shape NoUniqueness -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    arraySize Int
    0 TypeBase Shape NoUniqueness
    sv_t

            -- Generate an assertion and reshapes to ensure that sv and si' are the same
            -- size.
            SubExp
    cmp <-
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "write_cmp" (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
    I.int64) SubExp
    si_w SubExp
    sv_w
            Certs
    c <-
              [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert
                [Char]
    "write_cert"
                SubExp
    cmp
                ErrorMsg SubExp
    "length of index and value array does not match"
            Certs -> InternaliseM VName -> InternaliseM VName
    forall a. Certs -> InternaliseM a -> InternaliseM a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying Certs
    c (InternaliseM VName -> InternaliseM VName)
    -> InternaliseM VName -> InternaliseM VName
    forall a b. (a -> b) -> a -> b
    $
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp (VName -> [Char]
    baseString VName
    sv [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_write_sv") (Exp (Rep InternaliseM) -> InternaliseM VName)
    -> Exp (Rep InternaliseM) -> InternaliseM VName
    forall a b. (a -> b) -> a -> b
    $
                [SubExp] -> VName -> Exp SOACS
    forall rep. [SubExp] -> VName -> Exp rep
    shapeCoerce (Shape -> [SubExp]
    forall d. ShapeBase d -> [d]
    I.shapeDims (Shape -> Int -> Shape -> Shape
    reshapeOuter ([SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape [SubExp
    si_w]) Int
    1 Shape
    sv_shape)) VName
    sv

          [TypeBase Shape NoUniqueness]
    indexType <- (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType ([TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness])
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM [TypeBase Shape NoUniqueness]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    si'
          [VName]
    indexName <- (TypeBase Shape NoUniqueness -> InternaliseM VName)
    -> [TypeBase Shape NoUniqueness] -> 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 (\TypeBase Shape NoUniqueness
    _ -> [Char] -> InternaliseM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "write_index") [TypeBase Shape NoUniqueness]
    indexType
          [VName]
    valueNames <- Int -> InternaliseM VName -> InternaliseM [VName]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM ([TypeBase Shape NoUniqueness] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TypeBase Shape NoUniqueness]
    sv_ts) (InternaliseM VName -> InternaliseM [VName])
    -> InternaliseM VName -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ [Char] -> InternaliseM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "write_value"

          [TypeBase Shape NoUniqueness]
    sa_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    sas
          let bodyTypes :: [TypeBase Shape NoUniqueness]
    bodyTypes = [[TypeBase Shape NoUniqueness]] -> [TypeBase Shape NoUniqueness]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (Int
    -> [TypeBase Shape NoUniqueness] -> [[TypeBase Shape NoUniqueness]]
    forall a. Int -> a -> [a]
    replicate ([TypeBase Shape NoUniqueness] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TypeBase Shape NoUniqueness]
    sv_ts) [TypeBase Shape NoUniqueness]
    indexType) [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a. [a] -> [a] -> [a]
    ++ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. Int -> TypeBase Shape u -> TypeBase Shape u
    I.stripArray Int
    dim) [TypeBase Shape NoUniqueness]
    sa_ts
              paramTypes :: [TypeBase Shape NoUniqueness]
    paramTypes = [TypeBase Shape NoUniqueness]
    indexType [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a. Semigroup a => a -> a -> a
    <> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType [TypeBase Shape NoUniqueness]
    sv_ts
              bodyNames :: [VName]
    bodyNames = [VName]
    indexName [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    valueNames
              bodyParams :: [Param (TypeBase Shape NoUniqueness)]
    bodyParams = (VName
     -> TypeBase Shape NoUniqueness
     -> Param (TypeBase Shape NoUniqueness))
    -> [VName]
    -> [TypeBase Shape NoUniqueness]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (Attrs
    -> VName
    -> TypeBase Shape NoUniqueness
    -> Param (TypeBase Shape NoUniqueness)
    forall dec. Attrs -> VName -> dec -> Param dec
    I.Param Attrs
    forall a. Monoid a => a
    mempty) [VName]
    bodyNames [TypeBase Shape NoUniqueness]
    paramTypes

          -- This body is boring right now, as every input is exactly the output.
          -- But it can get funky later on if fused with something else.
          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 (TypeBase Shape NoUniqueness)] -> Scope SOACS
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams [Param (TypeBase Shape NoUniqueness)]
    bodyParams) (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
    -> (InternaliseM Result -> InternaliseM (Body SOACS))
    -> InternaliseM Result
    -> InternaliseM (Body SOACS)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
    InternaliseM Result -> InternaliseM (Body SOACS)
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (InternaliseM Result -> InternaliseM (Body SOACS))
    -> InternaliseM Result -> InternaliseM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $ do
            let outs :: [VName]
    outs = [[VName]] -> [VName]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (Int -> [VName] -> [[VName]]
    forall a. Int -> a -> [a]
    replicate ([VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    valueNames) [VName]
    indexName) [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    valueNames
            [SubExp]
    results <- [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    outs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
    -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    name ->
              [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "write_res" (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
    $ SubExp -> BasicOp
    I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    I.Var VName
    name
            ErrorMsg SubExp
    -> [TypeBase Shape NoUniqueness] -> Result -> InternaliseM Result
    ensureResultShape
              ErrorMsg SubExp
    "scatter value has wrong size"
              [TypeBase Shape NoUniqueness]
    bodyTypes
              ([SubExp] -> Result
    subExpsRes [SubExp]
    results)

          let lam :: Lambda SOACS
    lam =
                I.Lambda
                  { lambdaParams :: [LParam SOACS]
    I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
    [LParam SOACS]
    bodyParams,
                    lambdaReturnType :: [TypeBase Shape NoUniqueness]
    I.lambdaReturnType = [TypeBase Shape NoUniqueness]
    bodyTypes,
                    lambdaBody :: Body SOACS
    I.lambdaBody = Body SOACS
    body
                  }
              sivs :: [VName]
    sivs = [VName]
    si' [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    svs'

          let sa_ws :: [Shape]
    sa_ws = (TypeBase Shape NoUniqueness -> Shape)
    -> [TypeBase Shape NoUniqueness] -> [Shape]
    forall a b. (a -> b) -> [a] -> [b]
    map ([SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape ([SubExp] -> Shape)
    -> (TypeBase Shape NoUniqueness -> [SubExp])
    -> TypeBase Shape NoUniqueness
    -> Shape
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take Int
    dim ([SubExp] -> [SubExp])
    -> (TypeBase Shape NoUniqueness -> [SubExp])
    -> TypeBase Shape NoUniqueness
    -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims) [TypeBase Shape NoUniqueness]
    sa_ts
              spec :: [(Shape, Int, VName)]
    spec = [Shape] -> [Int] -> [VName] -> [(Shape, Int, VName)]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 [Shape]
    sa_ws (Int -> [Int]
    forall a. a -> [a]
    repeat Int
    1) [VName]
    sas
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
    -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall rep. Op rep -> Exp rep
    I.Op (Op (Rep InternaliseM) -> Exp (Rep InternaliseM))
    -> Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp
    -> [VName] -> [(Shape, Int, VName)] -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp
    -> [VName] -> [(Shape, Int, VName)] -> Lambda rep -> SOAC rep
    I.Scatter SubExp
    si_w [VName]
    sivs [(Shape, Int, VName)]
    spec Lambda SOACS
    lam

    flatIndexHelper :: String -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]
    flatIndexHelper :: [Char] -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
    flatIndexHelper [Char]
    desc Exp
    arr Exp
    offset [(Exp, Exp)]
    slices = do
      [VName]
    arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "arr" Exp
    arr
      SubExp
    offset' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "offset" Exp
    offset
      SubExp
    old_dim <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    I.arraysSize Int
    0 ([TypeBase Shape NoUniqueness] -> SubExp)
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arrs
      SubExp
    offset_inbounds_down <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "offset_inbounds_down" (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 (IntType -> CmpOp
    I.CmpUle IntType
    Int64) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExp
    offset'
      SubExp
    offset_inbounds_up <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "offset_inbounds_up" (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 (IntType -> CmpOp
    I.CmpUlt IntType
    Int64) SubExp
    offset' SubExp
    old_dim
      [(SubExp, SubExp)]
    slices' <-
        ((Exp, Exp) -> InternaliseM (SubExp, SubExp))
    -> [(Exp, Exp)] -> InternaliseM [(SubExp, 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
    n, Exp
    s) -> do
              SubExp
    n' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "n" Exp
    n
              SubExp
    s' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "s" Exp
    s
              (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    n', SubExp
    s')
          )
          [(Exp, Exp)]
    slices
      (SubExp
    min_bound, SubExp
    max_bound) <-
        ((SubExp, SubExp)
     -> (SubExp, SubExp) -> InternaliseM (SubExp, SubExp))
    -> (SubExp, SubExp)
    -> [(SubExp, SubExp)]
    -> InternaliseM (SubExp, SubExp)
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM
          ( \(SubExp
    lower, SubExp
    upper) (SubExp
    n, SubExp
    s) -> do
              SubExp
    n_m1 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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)
              SubExp
    spn <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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.Mul IntType
    Int64 Overflow
    I.OverflowUndef) SubExp
    n_m1 SubExp
    s

              SubExp
    span_and_lower <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "span_and_lower" (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.Add IntType
    Int64 Overflow
    I.OverflowUndef) SubExp
    spn SubExp
    lower
              SubExp
    span_and_upper <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "span_and_upper" (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.Add IntType
    Int64 Overflow
    I.OverflowUndef) SubExp
    spn SubExp
    upper

              SubExp
    lower' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "minimum" (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 -> BinOp
    I.UMin IntType
    Int64) SubExp
    span_and_lower SubExp
    lower
              SubExp
    upper' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "maximum" (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 -> BinOp
    I.UMax IntType
    Int64) SubExp
    span_and_upper SubExp
    upper

              (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    lower', SubExp
    upper')
          )
          (SubExp
    offset', SubExp
    offset')
          [(SubExp, SubExp)]
    slices'
      SubExp
    min_in_bounds <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "min_in_bounds" (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 (IntType -> CmpOp
    I.CmpUle IntType
    Int64) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExp
    min_bound
      SubExp
    max_in_bounds <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "max_in_bounds" (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 (IntType -> CmpOp
    I.CmpUlt IntType
    Int64) SubExp
    max_bound SubExp
    old_dim

      SubExp
    all_bounds <-
        (SubExp -> SubExp -> InternaliseM SubExp)
    -> SubExp -> [SubExp] -> InternaliseM SubExp
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM
          (\SubExp
    x SubExp
    y -> [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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)
          SubExp
    offset_inbounds_down
          [SubExp
    offset_inbounds_up, SubExp
    min_in_bounds, SubExp
    max_in_bounds]

      Certs
    c <-
        [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "bounds_cert" SubExp
    all_bounds (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 -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
    forall a b. (a -> b) -> a -> b
    $ Text
    "Flat slice out of bounds: " Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> Text
    forall a. Pretty a => a -> Text
    prettyText SubExp
    old_dim Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " and " Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [(SubExp, SubExp)] -> Text
    forall a. Pretty a => a -> Text
    prettyText [(SubExp, SubExp)]
    slices']
      let slice :: FlatSlice SubExp
    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'
      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
    $
        [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
    -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    arr' ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 :: String -> E.Exp -> E.Exp -> [E.Exp] -> E.Exp -> InternaliseM [SubExp]
    flatUpdateHelper :: [Char] -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
    flatUpdateHelper [Char]
    desc Exp
    arr1 Exp
    offset [Exp]
    slices Exp
    arr2 = do
      [VName]
    arrs1 <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "arr" Exp
    arr1
      SubExp
    offset' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "offset" Exp
    offset
      SubExp
    old_dim <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    I.arraysSize Int
    0 ([TypeBase Shape NoUniqueness] -> SubExp)
    -> InternaliseM [TypeBase Shape NoUniqueness]
    -> InternaliseM SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arrs1
      SubExp
    offset_inbounds_down <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "offset_inbounds_down" (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 (IntType -> CmpOp
    I.CmpUle IntType
    Int64) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExp
    offset'
      SubExp
    offset_inbounds_up <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "offset_inbounds_up" (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 (IntType -> CmpOp
    I.CmpUlt IntType
    Int64) SubExp
    offset' SubExp
    old_dim
      [VName]
    arrs2 <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "arr" Exp
    arr2
      [TypeBase Shape NoUniqueness]
    ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arrs2
      [(SubExp, SubExp)]
    slices' <-
        ((Exp, Int) -> InternaliseM (SubExp, SubExp))
    -> [(Exp, Int)] -> InternaliseM [(SubExp, 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
    s, Int
    i) -> do
              SubExp
    s' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "s" Exp
    s
              let n :: SubExp
    n = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    i [TypeBase Shape NoUniqueness]
    ts
              (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    n, SubExp
    s')
          )
          ([(Exp, Int)] -> InternaliseM [(SubExp, SubExp)])
    -> [(Exp, Int)] -> InternaliseM [(SubExp, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ [Exp] -> [Int] -> [(Exp, Int)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Exp]
    slices [Int
    0 ..]
      (SubExp
    min_bound, SubExp
    max_bound) <-
        ((SubExp, SubExp)
     -> (SubExp, SubExp) -> InternaliseM (SubExp, SubExp))
    -> (SubExp, SubExp)
    -> [(SubExp, SubExp)]
    -> InternaliseM (SubExp, SubExp)
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM
          ( \(SubExp
    lower, SubExp
    upper) (SubExp
    n, SubExp
    s) -> do
              SubExp
    n_m1 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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)
              SubExp
    spn <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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.Mul IntType
    Int64 Overflow
    I.OverflowUndef) SubExp
    n_m1 SubExp
    s

              SubExp
    span_and_lower <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "span_and_lower" (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.Add IntType
    Int64 Overflow
    I.OverflowUndef) SubExp
    spn SubExp
    lower
              SubExp
    span_and_upper <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "span_and_upper" (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.Add IntType
    Int64 Overflow
    I.OverflowUndef) SubExp
    spn SubExp
    upper

              SubExp
    lower' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "minimum" (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 -> BinOp
    I.UMin IntType
    Int64) SubExp
    span_and_lower SubExp
    lower
              SubExp
    upper' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "maximum" (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 -> BinOp
    I.UMax IntType
    Int64) SubExp
    span_and_upper SubExp
    upper

              (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    lower', SubExp
    upper')
          )
          (SubExp
    offset', SubExp
    offset')
          [(SubExp, SubExp)]
    slices'
      SubExp
    min_in_bounds <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "min_in_bounds" (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 (IntType -> CmpOp
    I.CmpUle IntType
    Int64) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExp
    min_bound
      SubExp
    max_in_bounds <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "max_in_bounds" (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 (IntType -> CmpOp
    I.CmpUlt IntType
    Int64) SubExp
    max_bound SubExp
    old_dim

      SubExp
    all_bounds <-
        (SubExp -> SubExp -> InternaliseM SubExp)
    -> SubExp -> [SubExp] -> InternaliseM SubExp
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM
          (\SubExp
    x SubExp
    y -> [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "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)
          SubExp
    offset_inbounds_down
          [SubExp
    offset_inbounds_up, SubExp
    min_in_bounds, SubExp
    max_in_bounds]

      Certs
    c <-
        [Char] -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
    assert [Char]
    "bounds_cert" SubExp
    all_bounds (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 -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
    forall a b. (a -> b) -> a -> b
    $ Text
    "Flat slice out of bounds: " Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> Text
    forall a. Pretty a => a -> Text
    prettyText SubExp
    old_dim Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " and " Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [(SubExp, SubExp)] -> Text
    forall a. Pretty a => a -> Text
    prettyText [(SubExp, SubExp)]
    slices']
      let slice :: FlatSlice SubExp
    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'
      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
    $
        [(VName, VName)]
    -> ((VName, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    arrs1 [VName]
    arrs2) (((VName, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
    -> ((VName, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    arr1', VName
    arr2') ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    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 ::
      String ->
      QualName VName ->
      [SubExp] ->
      InternaliseM [SubExp]
    funcall :: [Char] -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
    funcall [Char]
    desc (QualName [VName]
    _ VName
    fname) [SubExp]
    args = do
      ([VName]
    shapes, [DeclType]
    value_paramts, [Param DeclType]
    fun_params, [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
    rettype_fun) <- VName -> InternaliseM FunInfo
    lookupFunction VName
    fname
      [TypeBase Shape NoUniqueness]
    argts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    args

      [SubExp]
    shapeargs <- [VName]
    -> [FParam SOACS]
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [SubExp]
    argShapes [VName]
    shapes [Param DeclType]
    [FParam SOACS]
    fun_params [TypeBase Shape NoUniqueness]
    argts
      let diets :: [Diet]
    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.ObservePrim
              [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 Uniqueness -> Diet
    I.diet [DeclType]
    value_paramts
      [SubExp]
    args' <-
        ErrorMsg SubExp
    -> [VName]
    -> [TypeBase Shape NoUniqueness]
    -> [SubExp]
    -> InternaliseM [SubExp]
    forall u.
    Typed (TypeBase Shape u) =>
    ErrorMsg SubExp
    -> [VName]
    -> [TypeBase Shape u]
    -> [SubExp]
    -> InternaliseM [SubExp]
    ensureArgShapes
          ErrorMsg SubExp
    "function arguments of wrong shape"
          ((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]
    fun_params)
          ((Param DeclType -> TypeBase Shape NoUniqueness)
    -> [Param DeclType] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map Param DeclType -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
    I.paramType [Param DeclType]
    fun_params)
          ([SubExp]
    shapeargs [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [SubExp]
    args)
      [TypeBase Shape NoUniqueness]
    argts' <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    args'
      case [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
    rettype_fun ([(SubExp, TypeBase Shape NoUniqueness)]
     -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)])
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    -> [TypeBase Shape NoUniqueness]
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    args' [TypeBase Shape NoUniqueness]
    argts' of
        Maybe [(TypeBase ExtShape Uniqueness, 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 NoUniqueness] -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString [TypeBase Shape NoUniqueness]
    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 Uniqueness, RetAls)]
    ts -> do
          Safety
    safety <- InternaliseM Safety
    askSafety
          Attrs
    attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks InternaliseEnv -> Attrs
    envAttrs
          Attrs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
    attributing Attrs
    attrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
    -> (Exp SOACS -> InternaliseM [SubExp])
    -> Exp SOACS
    -> InternaliseM [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char] -> Exp SOACS -> InternaliseM [SubExp]
    letValExp' [Char]
    desc (Exp SOACS -> InternaliseM [SubExp])
    -> Exp SOACS -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $
            Name
    -> [(SubExp, Diet)]
    -> [(RetType SOACS, RetAls)]
    -> Safety
    -> Exp SOACS
    forall rep.
    Name
    -> [(SubExp, Diet)] -> [(RetType rep, RetAls)] -> Safety -> Exp rep
    I.Apply (VName -> Name
    internaliseFunName VName
    fname) ([SubExp] -> [Diet] -> [(SubExp, Diet)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    args' [Diet]
    diets) [(TypeBase ExtShape Uniqueness, RetAls)]
    [(RetType SOACS, RetAls)]
    ts Safety
    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 Uniqueness]
    ts = (Tree (TypeBase ExtShape Uniqueness)
     -> [TypeBase ExtShape Uniqueness])
    -> [Tree (TypeBase ExtShape Uniqueness)]
    -> [TypeBase ExtShape Uniqueness]
    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 Uniqueness)
    -> [TypeBase ExtShape Uniqueness]
    forall a. Free [] a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList ([Tree (TypeBase ExtShape Uniqueness)]
     -> [TypeBase ExtShape Uniqueness])
    -> [Tree (TypeBase ExtShape Uniqueness)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> a -> b
    $ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
    internaliseType (StructType -> [Tree (TypeBase ExtShape Uniqueness)])
    -> StructType -> [Tree (TypeBase ExtShape Uniqueness)]
    forall a b. (a -> b) -> a -> b
    $ StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    E.toStruct StructType
    ret
      [TypeBase Shape NoUniqueness]
    ses_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    subExpType [SubExp]
    ses

      let combine :: TypeBase ExtShape Uniqueness
    -> TypeBase Shape NoUniqueness -> Map VName SubExp
    combine TypeBase ExtShape Uniqueness
    t1 TypeBase Shape NoUniqueness
    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 Uniqueness -> [Ext SubExp]
    forall u. TypeBase ExtShape u -> [Ext SubExp]
    arrayExtDims TypeBase ExtShape Uniqueness
    t1) (TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims TypeBase Shape NoUniqueness
    t2)
          combine' :: Ext SubExp -> SubExp -> Map VName SubExp
    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

      [(VName, SubExp)]
    -> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ (Map VName SubExp -> [(VName, SubExp)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName SubExp -> [(VName, SubExp)])
    -> Map VName SubExp -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ [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
    $ (TypeBase ExtShape Uniqueness
     -> TypeBase Shape NoUniqueness -> Map VName SubExp)
    -> [TypeBase ExtShape Uniqueness]
    -> [TypeBase Shape NoUniqueness]
    -> [Map VName SubExp]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith TypeBase ExtShape Uniqueness
    -> TypeBase Shape NoUniqueness -> Map VName SubExp
    combine [TypeBase ExtShape Uniqueness]
    ts [TypeBase Shape NoUniqueness]
    ses_ts) (((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
    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
      Bool
    check <- (InternaliseEnv -> Bool) -> InternaliseM Bool
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks InternaliseEnv -> Bool
    envDoBoundsChecks
      Safety -> InternaliseM Safety
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Safety -> InternaliseM Safety) -> Safety -> InternaliseM Safety
    forall a b. (a -> b) -> a -> b
    $ if Bool
    check then Safety
    I.Safe else Safety
    I.Unsafe

    -- Implement partitioning using maps, scans and writes.
    partitionWithSOACS :: Int -> I.Lambda SOACS -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])
    partitionWithSOACS :: Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
    partitionWithSOACS Int
    k Lambda SOACS
    lam [VName]
    arrs = do
      [TypeBase Shape NoUniqueness]
    arr_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
    -> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
    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 -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    arrs
      let w :: SubExp
    w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    arr_ts
      [VName]
    classes_and_increments <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    "increments" (Exp (Rep InternaliseM) -> InternaliseM [VName])
    -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall rep. Op rep -> Exp rep
    I.Op (Op (Rep InternaliseM) -> Exp (Rep InternaliseM))
    -> Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    I.Screma SubExp
    w [VName]
    arrs (Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    mapSOAC Lambda SOACS
    lam)
      (VName
    classes, [VName]
    increments) <- case [VName]
    classes_and_increments of
        VName
    classes : [VName]
    increments -> (VName, [VName]) -> InternaliseM (VName, [VName])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName
    classes, Int -> [VName] -> [VName]
    forall a. Int -> [a] -> [a]
    take Int
    k [VName]
    increments)
        [VName]
    _ -> [Char] -> InternaliseM (VName, [VName])
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "partitionWithSOACS"

      [Param (TypeBase Shape NoUniqueness)]
    add_lam_x_params <-
        Int
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM Int
    k (InternaliseM (Param (TypeBase Shape NoUniqueness))
     -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall a b. (a -> b) -> a -> b
    $ [Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "x" (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64)
      [Param (TypeBase Shape NoUniqueness)]
    add_lam_y_params <-
        Int
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM Int
    k (InternaliseM (Param (TypeBase Shape NoUniqueness))
     -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall a b. (a -> b) -> a -> b
    $ [Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "y" (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64)
      Body SOACS
    add_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
    $
        Scope SOACS -> Builder SOACS Result -> Builder SOACS Result
    forall a.
    Scope SOACS
    -> BuilderT SOACS (State VNameSource) a
    -> BuilderT SOACS (State VNameSource) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS)
    -> [Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
    forall a b. (a -> b) -> a -> b
    $ [Param (TypeBase Shape NoUniqueness)]
    add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. [a] -> [a] -> [a]
    ++ [Param (TypeBase Shape NoUniqueness)]
    add_lam_y_params) (Builder SOACS Result -> Builder SOACS Result)
    -> Builder SOACS Result -> Builder SOACS Result
    forall a b. (a -> b) -> a -> b
    $
          ([SubExp] -> Result)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    -> Builder SOACS Result
    forall a b.
    (a -> b)
    -> BuilderT SOACS (State VNameSource) a
    -> BuilderT SOACS (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [SubExp] -> Result
    subExpsRes (BuilderT SOACS (State VNameSource) [SubExp]
     -> Builder SOACS Result)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    -> Builder SOACS Result
    forall a b. (a -> b) -> a -> b
    $
            [(Param (TypeBase Shape NoUniqueness),
      Param (TypeBase Shape NoUniqueness))]
    -> ((Param (TypeBase Shape NoUniqueness),
         Param (TypeBase Shape NoUniqueness))
        -> BuilderT SOACS (State VNameSource) SubExp)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [(Param (TypeBase Shape NoUniqueness),
         Param (TypeBase Shape NoUniqueness))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param (TypeBase Shape NoUniqueness)]
    add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
    add_lam_y_params) (((Param (TypeBase Shape NoUniqueness),
       Param (TypeBase Shape NoUniqueness))
      -> BuilderT SOACS (State VNameSource) SubExp)
     -> BuilderT SOACS (State VNameSource) [SubExp])
    -> ((Param (TypeBase Shape NoUniqueness),
         Param (TypeBase Shape NoUniqueness))
        -> BuilderT SOACS (State VNameSource) SubExp)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \(Param (TypeBase Shape NoUniqueness)
    x, Param (TypeBase Shape NoUniqueness)
    y) ->
              [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "z" (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall a b. (a -> b) -> a -> b
    $
                BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $
                  BinOp -> SubExp -> SubExp -> BasicOp
    I.BinOp
                    (IntType -> Overflow -> BinOp
    I.Add IntType
    Int64 Overflow
    I.OverflowUndef)
                    (VName -> SubExp
    I.Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName Param (TypeBase Shape NoUniqueness)
    x)
                    (VName -> SubExp
    I.Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName Param (TypeBase Shape NoUniqueness)
    y)
      let add_lam :: Lambda SOACS
    add_lam =
            I.Lambda
              { lambdaBody :: Body SOACS
    I.lambdaBody = Body SOACS
    add_lam_body,
                lambdaParams :: [LParam SOACS]
    I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
    add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. [a] -> [a] -> [a]
    ++ [Param (TypeBase Shape NoUniqueness)]
    add_lam_y_params,
                lambdaReturnType :: [TypeBase Shape NoUniqueness]
    I.lambdaReturnType = Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
    forall a. Int -> a -> [a]
    replicate Int
    k (TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness])
    -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64
              }
          nes :: [SubExp]
    nes = Int -> SubExp -> [SubExp]
    forall a. Int -> a -> [a]
    replicate ([VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    increments) (SubExp -> [SubExp]) -> SubExp -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0

      ScremaForm SOACS
    scan <- [Scan SOACS] -> InternaliseM (ScremaForm SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Scan rep] -> m (ScremaForm rep)
    I.scanSOAC [Lambda SOACS -> [SubExp] -> Scan SOACS
    forall rep. Lambda rep -> [SubExp] -> Scan rep
    I.Scan Lambda SOACS
    add_lam [SubExp]
    nes]
      [VName]
    all_offsets <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    "offsets" (Exp (Rep InternaliseM) -> InternaliseM [VName])
    -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall rep. Op rep -> Exp rep
    I.Op (Op (Rep InternaliseM) -> Exp (Rep InternaliseM))
    -> Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    I.Screma SubExp
    w [VName]
    increments ScremaForm SOACS
    scan

      -- We have the offsets for each of the partitions, but we also need
      -- the total sizes, which are the last elements in the offests.  We
      -- just have to be careful in case the array is empty.
      SubExp
    last_index <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "last_index" (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
    OverflowUndef) SubExp
    w (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64
    1 :: Int64)
      let nonempty_body :: InternaliseM (Body SOACS)
    nonempty_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
    $
            ([SubExp] -> Result)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    -> Builder SOACS Result
    forall a b.
    (a -> b)
    -> BuilderT SOACS (State VNameSource) a
    -> BuilderT SOACS (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [SubExp] -> Result
    subExpsRes (BuilderT SOACS (State VNameSource) [SubExp]
     -> Builder SOACS Result)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    -> Builder SOACS Result
    forall a b. (a -> b) -> a -> b
    $
              [VName]
    -> (VName -> BuilderT SOACS (State VNameSource) SubExp)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    all_offsets ((VName -> BuilderT SOACS (State VNameSource) SubExp)
     -> BuilderT SOACS (State VNameSource) [SubExp])
    -> (VName -> BuilderT SOACS (State VNameSource) SubExp)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    offset_array ->
                [Char]
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "last_offset" (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) SubExp
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ VName -> Slice SubExp -> BasicOp
    I.Index VName
    offset_array (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    I.DimFix SubExp
    last_index]
          empty_body :: InternaliseM (Body (Rep InternaliseM))
    empty_body = [SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM ([SubExp] -> InternaliseM (Body (Rep InternaliseM)))
    -> [SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $ Int -> SubExp -> [SubExp]
    forall a. Int -> a -> [a]
    replicate Int
    k (SubExp -> [SubExp]) -> SubExp -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64
    0 :: Int64)
      SubExp
    is_empty <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "is_empty" (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
    CmpEq PrimType
    int64) SubExp
    w (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64
    0 :: Int64)
      [VName]
    sizes <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    "partition_size" (Exp SOACS -> InternaliseM [VName])
    -> InternaliseM (Exp SOACS) -> InternaliseM [VName]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< InternaliseM (Exp (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    is_empty) InternaliseM (Body (Rep InternaliseM))
    empty_body InternaliseM (Body (Rep InternaliseM))
    InternaliseM (Body SOACS)
    nonempty_body

      -- The total size of all partitions must necessarily be equal to the
      -- size of the input array.

      -- Create scratch arrays for the result.
      [VName]
    blanks <- [TypeBase Shape NoUniqueness]
    -> (TypeBase Shape NoUniqueness -> InternaliseM VName)
    -> InternaliseM [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [TypeBase Shape NoUniqueness]
    arr_ts ((TypeBase Shape NoUniqueness -> InternaliseM VName)
     -> InternaliseM [VName])
    -> (TypeBase Shape NoUniqueness -> InternaliseM VName)
    -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $ \TypeBase Shape NoUniqueness
    arr_t ->
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "partition_dest" (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
    $
            PrimType -> [SubExp] -> BasicOp
    Scratch (TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    I.elemType TypeBase Shape NoUniqueness
    arr_t) (SubExp
    w SubExp -> [SubExp] -> [SubExp]
    forall a. a -> [a] -> [a]
    : Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    drop Int
    1 (TypeBase Shape NoUniqueness -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    I.arrayDims TypeBase Shape NoUniqueness
    arr_t))

      -- Now write into the result.
      Lambda SOACS
    write_lam <- do
        Param (TypeBase Shape NoUniqueness)
    c_param <- [Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "c" (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64)
        [Param (TypeBase Shape NoUniqueness)]
    offset_params <- Int
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM Int
    k (InternaliseM (Param (TypeBase Shape NoUniqueness))
     -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    forall a b. (a -> b) -> a -> b
    $ [Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "offset" (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64)
        [Param (TypeBase Shape NoUniqueness)]
    value_params <- (TypeBase Shape NoUniqueness
     -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
    -> [TypeBase Shape NoUniqueness]
    -> InternaliseM [Param (TypeBase Shape NoUniqueness)]
    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 ([Char]
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam [Char]
    "v" (TypeBase Shape NoUniqueness
     -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
    -> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    I.rowType) [TypeBase Shape NoUniqueness]
    arr_ts
        (SubExp
    offset, Stms SOACS
    offset_stms) <-
          InternaliseM SubExp
    -> InternaliseM (SubExp, 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 SubExp
     -> InternaliseM (SubExp, Stms (Rep InternaliseM)))
    -> InternaliseM SubExp
    -> InternaliseM (SubExp, Stms (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $
            [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
    mkOffsetLambdaBody
              ((VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var [VName]
    sizes)
              (VName -> SubExp
    I.Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName Param (TypeBase Shape NoUniqueness)
    c_param)
              Int
    0
              [Param (TypeBase Shape NoUniqueness)]
    [LParam SOACS]
    offset_params
        Lambda SOACS -> InternaliseM (Lambda SOACS)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
          I.Lambda
            { lambdaParams :: [LParam SOACS]
    I.lambdaParams = Param (TypeBase Shape NoUniqueness)
    c_param Param (TypeBase Shape NoUniqueness)
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. a -> [a] -> [a]
    : [Param (TypeBase Shape NoUniqueness)]
    offset_params [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. [a] -> [a] -> [a]
    ++ [Param (TypeBase Shape NoUniqueness)]
    value_params,
              lambdaReturnType :: [TypeBase Shape NoUniqueness]
    I.lambdaReturnType =
                Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
    forall a. Int -> a -> [a]
    replicate ([TypeBase Shape NoUniqueness] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TypeBase Shape NoUniqueness]
    arr_ts) (PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64)
                  [TypeBase Shape NoUniqueness]
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a. [a] -> [a] -> [a]
    ++ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
    -> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
    forall u. TypeBase Shape u -> TypeBase Shape u
    I.rowType [TypeBase Shape NoUniqueness]
    arr_ts,
              lambdaBody :: Body SOACS
    I.lambdaBody =
                Stms SOACS -> Result -> Body SOACS
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms SOACS
    offset_stms (Result -> Body SOACS) -> Result -> Body SOACS
    forall a b. (a -> b) -> a -> b
    $
                  Int -> SubExpRes -> Result
    forall a. Int -> a -> [a]
    replicate ([TypeBase Shape NoUniqueness] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TypeBase Shape NoUniqueness]
    arr_ts) (SubExp -> SubExpRes
    subExpRes SubExp
    offset)
                    Result -> Result -> Result
    forall a. [a] -> [a] -> [a]
    ++ [VName] -> Result
    I.varsRes ((Param (TypeBase Shape NoUniqueness) -> VName)
    -> [Param (TypeBase Shape NoUniqueness)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName [Param (TypeBase Shape NoUniqueness)]
    value_params)
            }
      let spec :: [(Shape, Int, VName)]
    spec = [Shape] -> [Int] -> [VName] -> [(Shape, Int, VName)]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 (Shape -> [Shape]
    forall a. a -> [a]
    repeat (Shape -> [Shape]) -> Shape -> [Shape]
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    I.Shape [SubExp
    w]) (Int -> [Int]
    forall a. a -> [a]
    repeat Int
    1) [VName]
    blanks
      [VName]
    results <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    "partition_res" (Exp SOACS -> InternaliseM [VName])
    -> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op SOACS -> Exp SOACS
    SOAC SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    I.Op (SOAC SOACS -> InternaliseM [VName])
    -> SOAC SOACS -> InternaliseM [VName]
    forall a b. (a -> b) -> a -> b
    $
          SubExp
    -> [VName] -> [(Shape, Int, VName)] -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp
    -> [VName] -> [(Shape, Int, VName)] -> Lambda rep -> SOAC rep
    I.Scatter SubExp
    w (VName
    classes VName -> [VName] -> [VName]
    forall a. a -> [a] -> [a]
    : [VName]
    all_offsets [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    arrs) [(Shape, Int, VName)]
    spec Lambda SOACS
    write_lam
      SubExp
    sizes' <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "partition_sizes" (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
    $
            [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
    I.ArrayLit ((VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var [VName]
    sizes) (TypeBase Shape NoUniqueness -> BasicOp)
    -> TypeBase Shape NoUniqueness -> BasicOp
    forall a b. (a -> b) -> a -> b
    $
              PrimType -> TypeBase Shape NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    I.Prim PrimType
    int64
      ([SubExp], [SubExp]) -> InternaliseM ([SubExp], [SubExp])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ((VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    I.Var [VName]
    results, [SubExp
    sizes'])
      where
        mkOffsetLambdaBody ::
          [SubExp] ->
          SubExp ->
          Int ->
          [I.LParam SOACS] ->
          InternaliseM SubExp
        mkOffsetLambdaBody :: [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
    mkOffsetLambdaBody [SubExp]
    _ SubExp
    _ Int
    _ [] =
          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
    $ Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (-Int64
    1 :: Int64)
        mkOffsetLambdaBody [SubExp]
    sizes SubExp
    c Int
    i (LParam SOACS
    p : [LParam SOACS]
    ps) = do
          SubExp
    is_this_one <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "is_this_one" (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
    CmpEq PrimType
    int64) SubExp
    c (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
          SubExp
    next_one <- [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
    mkOffsetLambdaBody [SubExp]
    sizes SubExp
    c (Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1) [LParam SOACS]
    ps
          SubExp
    this_one <-
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "this_offset"
              (Exp SOACS -> InternaliseM SubExp)
    -> InternaliseM (Exp SOACS) -> InternaliseM SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< BinOp
    -> SubExp -> [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    foldBinOp
                (IntType -> Overflow -> BinOp
    Add IntType
    Int64 Overflow
    OverflowUndef)
                (Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (-Int64
    1 :: Int64))
                (VName -> SubExp
    I.Var (Param (TypeBase Shape NoUniqueness) -> VName
    forall dec. Param dec -> VName
    I.paramName Param (TypeBase Shape NoUniqueness)
    LParam SOACS
    p) SubExp -> [SubExp] -> [SubExp]
    forall a. a -> [a] -> [a]
    : Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take Int
    i [SubExp]
    sizes)
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "total_res"
            (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 (Body (Rep InternaliseM))
    -> InternaliseM (Body (Rep InternaliseM))
    -> InternaliseM (Exp (Rep InternaliseM))
    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 -> InternaliseM (Exp (Rep InternaliseM))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    is_this_one)
              ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    this_one])
              ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [SubExp
    next_one])

    sizeExpForError :: E.Size -> InternaliseM [ErrorMsgPart SubExp]
    sizeExpForError :: Exp -> InternaliseM [ErrorMsgPart SubExp]
    sizeExpForError Exp
    e
      | Exp
    e Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    == Exp
    anySize = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [ErrorMsgPart SubExp
    "[]"]
      | Bool
    otherwise = do
          SubExp
    e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "size" Exp
    e
          [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [ErrorMsgPart SubExp
    "[", PrimType -> SubExp -> ErrorMsgPart SubExp
    forall a. PrimType -> a -> ErrorMsgPart a
    ErrorVal PrimType
    int64 SubExp
    e', ErrorMsgPart SubExp
    "]"]

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

    errorShape :: [a] -> ErrorMsg a
    errorShape :: forall a. [a] -> ErrorMsg a
    errorShape [a]
    dims =
      ErrorMsg a
    "["
        ErrorMsg a -> ErrorMsg a -> ErrorMsg a
    forall a. Semigroup a => a -> a -> a
    <> [ErrorMsg a] -> ErrorMsg a
    forall a. Monoid a => [a] -> a
    mconcat (ErrorMsg a -> [ErrorMsg a] -> [ErrorMsg a]
    forall a. a -> [a] -> [a]
    intersperse ErrorMsg a
    "][" ([ErrorMsg a] -> [ErrorMsg a]) -> [ErrorMsg a] -> [ErrorMsg a]
    forall a b. (a -> b) -> a -> b
    $ (a -> ErrorMsg a) -> [a] -> [ErrorMsg a]
    forall a b. (a -> b) -> [a] -> [b]
    map ([ErrorMsgPart a] -> ErrorMsg a
    forall a. [ErrorMsgPart a] -> ErrorMsg a
    ErrorMsg ([ErrorMsgPart a] -> ErrorMsg a)
    -> (a -> [ErrorMsgPart a]) -> a -> ErrorMsg a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ErrorMsgPart a -> [ErrorMsgPart a]
    forall a. a -> [a]
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (ErrorMsgPart a -> [ErrorMsgPart a])
    -> (a -> ErrorMsgPart a) -> a -> [ErrorMsgPart a]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PrimType -> a -> ErrorMsgPart a
    forall a. PrimType -> a -> ErrorMsgPart a
    ErrorVal PrimType
    int64) [a]
    dims)
        ErrorMsg a -> ErrorMsg a -> ErrorMsg a
    forall a. Semigroup a => a -> a -> a
    <> ErrorMsg a
    "]"