    {-# 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
      (opaques, consts, 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)
      I.renameProg $ I.Prog opaques consts 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', 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
          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
          (rettype', retals) <-
            first zeroExts . unzip . internaliseReturnType (map (fmap paramDeclType) all_params) rettype
              <$> mapM subExpType body_res

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

          body_res' <-
            ensureResultExtShape msg (map I.fromDecl rettype') $ subExpsRes body_res
          let 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')
          pure
            ( body_res',
              replicate num_ctx (I.Prim int64, mempty)
                ++ zip rettype' (map (shiftRetAls num_ctx) retals)
            )

        attrs' <- internaliseAttrs attrs

        let fd =
              Maybe EntryPoint
    -> Attrs
    -> Name
    -> [(RetType SOACS, RetAls)]
    -> [FParam SOACS]
    -> Body SOACS
    -> FunDef SOACS
    forall rep.
    Maybe EntryPoint
    -> Attrs
    -> Name
    -> [(RetType rep, RetAls)]
    -> [FParam rep]
    -> Body rep
    -> FunDef rep
    I.FunDef
                Maybe EntryPoint
    forall a. Maybe a
    Nothing
                Attrs
    attrs'
                (VName -> Name
    internaliseFunName VName
    fname)
                [(TypeBase ExtShape 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 null params'
          then bindConstant fname fd
          else
            bindFunction
              fname
              fd
              ( shapenames,
                map declTypeOf $ foldMap (foldMap toList) params',
                foldMap toList all_params,
                fmap (`zip` map snd rettype')
                  . applyRetType (map fst rettype') (foldMap toList 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

        (entry_body, 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_const <- VName -> InternaliseM (Maybe [SubExp])
    lookupConst VName
    ofname
          vals <- case maybe_const of
            Just [SubExp]
    ses ->
              [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp]
    ses
            Maybe [SubExp]
    Nothing ->
              [Char] -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
    funcall [Char]
    "entry_result" (VName -> QualName VName
    forall v. v -> QualName v
    E.qualName VName
    ofname) [SubExp]
    args
          ctx <-
            extractShapeContext (zeroExts $ concat entry_rettype)
              <$> mapM (fmap I.arrayDims . subExpType) vals
          pure (subExpsRes $ ctx ++ vals, map (const (I.Prim int64, mempty)) ctx)

        attrs' <- internaliseAttrs attrs
        let num_ctx = [(TypeBase ExtShape Uniqueness, RetAls)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(TypeBase ExtShape Uniqueness, RetAls)]
    ctx_ts
        addFunDef $
          I.FunDef
            (Just entry')
            attrs'
            ("entry_" <> baseName ofname)
            ( ctx_ts
                ++ zip
                  (zeroExts (concat entry_rettype))
                  (map (shiftRetAls num_ctx) $ concat retals)
            )
            (shapeparams ++ foldMap (foldMap toList) params')
            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
      ((res, a), 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) <$> mkBodyM stms 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
      e_t <- Exp SOACS -> InternaliseM [ExtType]
    forall rep (m :: * -> *).
    (HasScope rep m, TypedOp (OpC rep)) =>
    Exp rep -> m [ExtType]
    expExtType Exp SOACS
    e
      names <- replicateM (length e_t) $ newVName name
      letBindNames names e
      let ctx = [ExtType] -> Set Int
    forall u. [TypeBase ExtShape u] -> Set Int
    shapeContext [ExtType]
    e_t
      pure $ map fst $ filter ((`S.notMember` ctx) . snd) $ zip names [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
      vs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "indexed" Exp
    e
      dims <- case vs of
        [] -> [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [] -- Will this happen?
        VName
    v : [VName]
    _ -> TypeBase Shape 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
      (idxs', cs) <- internaliseSlice dims idxs
      let index VName
    v = do
            v_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    v
            pure $ I.BasicOp $ I.Index v $ fullSlice v_t idxs'
      certifying cs $ mapM (letSubExp desc <=< index) vs
    internaliseAppExp [Char]
    desc AppRes
    _ (E.Range Exp
    start Maybe Exp
    maybe_second Inclusiveness Exp
    end SrcLoc
    _) = do
      start' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "range_start" Exp
    start
      end' <- internaliseExp1 "range_end" $ case end of
        DownToExclusive Exp
    e -> Exp
    e
        ToInclusive Exp
    e -> Exp
    e
        UpToExclusive Exp
    e -> Exp
    e
      maybe_second' <-
        traverse (internaliseExp1 "range_second") maybe_second

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

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

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

      (step, step_zero) <- case maybe_second' of
        Just SubExp
    second' -> do
          subtracted_step <-
            [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'
          step_zero <- letSubExp "step_zero" $ I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) start' second'
          pure (subtracted_step, step_zero)
        Maybe SubExp
    Nothing ->
          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    default_step, Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    False)

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

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

      (distance, step_wrong_dir, bounds_invalid) <- case end of
        DownToExclusive {} -> do
          step_wrong_dir <-
            [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
          distance <-
            letSubExp "distance" $
              I.BasicOp $
                I.BinOp (Sub it I.OverflowWrap) start' end'
          distance_i64 <- asIntS Int64 distance
          pure (distance_i64, step_wrong_dir, bounds_invalid_downwards)
        UpToExclusive {} -> do
          step_wrong_dir <-
            [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
          distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start'
          distance_i64 <- asIntS Int64 distance
          pure (distance_i64, step_wrong_dir, bounds_invalid_upwards)
        ToInclusive {} -> do
          downwards <-
            [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
          distance_downwards_exclusive <-
            letSubExp "distance_downwards_exclusive" $
              I.BasicOp $
                I.BinOp (Sub it I.OverflowWrap) start' end'
          distance_upwards_exclusive <-
            letSubExp "distance_upwards_exclusive" $
              I.BasicOp $
                I.BinOp (Sub it I.OverflowWrap) end' start'

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

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

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

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

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

      se <- letSubExp desc (I.BasicOp $ I.Iota num_elems start' step it)
      pure [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]
                  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] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap [SubExp] -> [(SubExp, Diet)]
    forall {a}. [a] -> [(a, Diet)]
    tag [[SubExp]]
    args'
                  letValExp' desc $ I.Apply fname args'' [(I.Prim rettype, mempty)] Safe
              | Bool
    otherwise -> do
                  args' <- [[SubExp]] -> [SubExp]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat ([[SubExp]] -> [SubExp])
    -> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [[SubExp]] -> [[SubExp]]
    forall a. [a] -> [a]
    reverse ([[SubExp]] -> [SubExp])
    -> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
    -> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM ([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)
                  funcall desc qfname 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
      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
      ((loopbody', (form', shapepat, mergepat', mergeinit')), initstms) <-
        collectStms $ handleForm ses form

      addStms initstms
      mergeinit_ts' <- mapM subExpType mergeinit'

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

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

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

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

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

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

        loopBody :: InternaliseM [SubExp]
    loopBody = (InternaliseEnv -> InternaliseEnv)
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a.
    (InternaliseEnv -> InternaliseEnv)
    -> InternaliseM a -> InternaliseM a
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
    local InternaliseEnv -> InternaliseEnv
    noLoopAttrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ [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)]
            ses <- InternaliseM [SubExp]
    loopBody
            sets <- mapM subExpType ses
            shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets
            pure
              ( subExpsRes $ shapeargs ++ ses,
                ( form',
                  shapepat,
                  mergepat',
                  mergeinit
                )
              )

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

          i <- newVName "i"

          ts <- mapM subExpType mergeinit
          bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
    shapepat [FParam SOACS]
    mergepat' ->
            [PatBase Info VName ParamType]
    -> [TypeBase Shape 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
          num_iterations' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "upper_bound" Exp
    num_iterations
          num_iterations_t <- I.subExpType num_iterations'
          it <- case num_iterations_t of
            I.Prim (IntType IntType
    it) -> IntType -> InternaliseM IntType
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure IntType
    it
            TypeBase Shape NoUniqueness
    _ -> [Char] -> InternaliseM IntType
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "internaliseExp Loop: invalid type"

          ts <- mapM subExpType mergeinit
          bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
    shapepat [FParam SOACS]
    mergepat' ->
            [Param DeclType]
    -> [Param DeclType]
    -> [SubExp]
    -> VName
    -> [(Param (TypeBase Shape 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
          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
          bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
    shapepat [FParam SOACS]
    mergepat' -> do
            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.
            shapeinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts

            (loop_initial_cond, init_loop_cond_stms) <- collectStms $ do
              forM_ (zip shapepat shapeinit) $ \(Param DeclType
    p, SubExp
    se) ->
                [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName Param DeclType
    p] (Exp (Rep InternaliseM) -> InternaliseM ())
    -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep InternaliseM)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep InternaliseM))
    -> BasicOp -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp SubExp
    se
              forM_ (zip mergepat' mergeinit) $ \(Param DeclType
    p, SubExp
    se) ->
                Bool -> InternaliseM () -> InternaliseM ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (SubExp
    se SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName -> SubExp
    I.Var (Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName Param DeclType
    p)) (InternaliseM () -> InternaliseM ())
    -> InternaliseM () -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $
                  [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName Param DeclType
    p] (Exp (Rep InternaliseM) -> InternaliseM ())
    -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $
                    case SubExp
    se of
                      I.Var VName
    v
                        | Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ TypeBase Shape 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).
              (cond_stms, cond') <-
                uncurry (flip renameStmsWith)
                  =<< collectStms (internaliseExp1 "loop_cond" cond)
              addStms cond_stms
              pure cond'

            addStms init_loop_cond_stms

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

              -- Careful not to clobber anything.
              loop_end_cond_body <- renameBody <=< buildBody_ $ do
                forM_ (zip shapepat shapeargs) $ \(Param DeclType
    p, SubExp
    se) ->
                  Bool -> InternaliseM () -> InternaliseM ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (SubExp
    se SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName -> SubExp
    I.Var (Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName Param DeclType
    p)) (InternaliseM () -> InternaliseM ())
    -> InternaliseM () -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $
                    [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [Param DeclType -> VName
    forall dec. Param dec -> VName
    I.paramName Param DeclType
    p] (Exp (Rep InternaliseM) -> InternaliseM ())
    -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $
                      BasicOp -> Exp (Rep InternaliseM)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep InternaliseM))
    -> BasicOp -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $
                        SubExp -> BasicOp
    SubExp SubExp
    se
                forM_ (zip mergepat' ses) $ \(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
                subExpsRes <$> internaliseExp "loop_cond" cond
              loop_end_cond <- bodyBind loop_end_cond_body

              pure
                ( subExpsRes shapeargs ++ loop_end_cond ++ subExpsRes ses,
                  ( I.WhileLoop $ I.paramName loop_while,
                    shapepat,
                    loop_while : mergepat',
                    loop_initial_cond : mergeinit
                  )
                )
    internaliseAppExp [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
      ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp ([Char]
    desc [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_scrutinee") Exp
    e
      cs <- mapM (onCase ses) orig_cs
      case NE.uncons cs of
        (I.Case [Maybe PrimValue]
    _ InternaliseM (Body SOACS)
    body, Maybe (NonEmpty (Case (InternaliseM (Body SOACS))))
    Nothing) ->
          (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
          (cmps, pertinent) <- PatBase Info VName StructType
    -> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
    generateCond PatBase Info VName StructType
    p [SubExp]
    ses
          pure . I.Case cmps $
            internalisePat' [] p pertinent $
              internaliseBody "case" case_e
    internaliseAppExp [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 (ZonkAny 0) -> Text
    forall a. Doc a -> Text
    docText (Doc (ZonkAny 0) -> Text) -> Doc (ZonkAny 0) -> Text
    forall a b. (a -> b) -> a -> b
    $ Doc (ZonkAny 0)
    "Reached hole of type: " Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
    forall a. Semigroup a => a -> a -> a
    <> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
    forall ann. Doc ann -> Doc ann
    align (StructType -> Doc (ZonkAny 0)
    forall a ann. Pretty a => a -> Doc ann
    forall ann. StructType -> Doc ann
    pretty StructType
    t)
          ts :: [TypeBase ExtShape 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)
      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 mapM hasStaticShape 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
      subst <- VName -> InternaliseM (Maybe [SubExp])
    lookupSubst VName
    name
      case subst of
        Just [SubExp]
    substs -> [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp]
    substs
        Maybe [SubExp]
    Nothing -> [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [VName -> SubExp
    I.Var VName
    name]
    internaliseExp [Char]
    desc (E.AppExp AppExp
    e (Info AppRes
    appres)) = do
      ses <- AppExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a b. Located a => a -> InternaliseM b -> InternaliseM b
    locating AppExp
    e (InternaliseM [SubExp] -> InternaliseM [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall a b. (a -> b) -> a -> b
    $ [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
    internaliseAppExp [Char]
    desc AppRes
    appres AppExp
    e
      bindExtSizes appres ses
      pure 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] -> [Int] -> Bool
    forall a. Eq a => a -> a -> Bool
    ==) ([Int] -> Bool)
    -> (([Int], [Exp]) -> [Int]) -> ([Int], [Exp]) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([Int], [Exp]) -> [Int]
    forall a b. (a, b) -> a
    fst) [([Int], [Exp])]
    es',
        Just StructType
    basetype <- Int -> StructType -> Maybe StructType
    forall dim 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
          flat_arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "flat_literal" Exp
    flat_lit
          forM flat_arrs $ \VName
    flat_arr -> do
            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 -> 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
            letSubExp desc $ I.BasicOp $ I.Reshape flat_arr (reshapeAll (I.arrayShape flat_arr_t) new_shape')
      | Bool
    otherwise = do
          es' <- (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM ([Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "arr_elem") [Exp]
    es
          let 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

          rowtypes <-
            case mapM (fmap rowType . hasStaticShape . I.fromDecl) 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]
    ks TypeBase Shape NoUniqueness
    rt = do
                ks' <-
                  (SubExp -> InternaliseM SubExp)
    -> [SubExp] -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM
                    ( ErrorMsg SubExp
    -> TypeBase Shape 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
                pure $ I.BasicOp $ I.ArrayLit ks' rt

          mapM (letSubExp desc)
            =<< if null es'
              then mapM (arraylit []) rowtypes
              else zipWithM arraylit (transpose es') rowtypes
      where
        isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
        isArrayLiteral :: Exp -> Maybe ([Int], [Exp])
    isArrayLiteral (E.ArrayLit [Exp]
    inner_es Info StructType
    _ SrcLoc
    _) = do
          (eshape, e) : inner_es' <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM Exp -> Maybe ([Int], [Exp])
    isArrayLiteral [Exp]
    inner_es
          guard $ all ((eshape ==) . fst) inner_es'
          pure (length inner_es : eshape, e ++ concatMap snd inner_es')
        isArrayLiteral Exp
    e =
          ([Int], [Exp]) -> Maybe ([Int], [Exp])
    forall a. a -> Maybe a
    Just ([], [Exp
    e])
    internaliseExp [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
      ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    e
      ts <- internaliseCoerceType (E.toStruct et) <$> mapM subExpType ses
      dt' <- typeExpForError $ toStruct et
      forM (zip ses ts) $ \(SubExp
    e', TypeBase ExtShape Uniqueness
    t') -> do
        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
    "Value of (desugared) shape ["]
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ ErrorMsgPart SubExp
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. a -> [a] -> [a]
    intersperse ErrorMsgPart SubExp
    "][" ((SubExp -> ErrorMsgPart SubExp)
    -> [SubExp] -> [ErrorMsgPart SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (PrimType -> SubExp -> ErrorMsgPart SubExp
    forall a. PrimType -> a -> ErrorMsgPart a
    ErrorVal PrimType
    int64) [SubExp]
    dims)
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [ErrorMsgPart SubExp
    "] cannot match shape of type \""]
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [ErrorMsgPart SubExp]
    dt'
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [ErrorMsgPart SubExp
    "\"."]
        ensureExtShape (errorMsg parts) (I.fromDecl t') desc 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
      e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "negate_arg" Exp
    e
      et <- subExpType e'
      case 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
      e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "not_arg" Exp
    e
      et <- subExpType e'
      case 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
      ves <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "lw_val" Exp
    ve
      srcs <- internaliseExpToVars "src" src
      (src_dims, ve_dims) <- case (srcs, 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?
      (idxs', cs) <- internaliseSlice src_dims slice
      let 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 = [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    src_dims'
          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
    sname SubExp
    ve' = do
            sname_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    sname
            let full_slice = TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
    fullSlice TypeBase Shape NoUniqueness
    sname_t [DimIndex SubExp]
    idxs'
                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
            ve'' <-
              ensureShape errormsg rowtype "lw_val_correct_shape" ve'
            letInPlace desc sname full_slice $ BasicOp $ SubExp ve''
      certifying cs $ map I.Var <$> zipWithM comb srcs 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
      src' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    src
      ve' <- internaliseExp desc ve
      replace (E.typeOf src) fields ve' 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'
              src'' <- TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
    replace TypeBase Exp als
    t [Name]
    fs [a]
    ve' [a]
    to_update
              pure $ bef ++ src'' ++ 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' <- AttrInfo VName -> InternaliseM Attr
    internaliseAttr AttrInfo VName
    attr
      e' <- local (f attr') $ internaliseExp desc e
      case attr' of
        Attr
    "trace" ->
          Text -> [SubExp] -> InternaliseM [SubExp]
    traceRes ([Char] -> Text
    T.pack ([Char] -> Text) -> [Char] -> Text
    forall a b. (a -> b) -> a -> b
    $ SrcLoc -> [Char]
    forall a. Located a => a -> [Char]
    locStr SrcLoc
    loc) [SubExp]
    e'
        I.AttrComp Name
    "trace" [I.AttrName Name
    tag] ->
          Text -> [SubExp] -> InternaliseM [SubExp]
    traceRes (Name -> Text
    nameToText Name
    tag) [SubExp]
    e'
        Attr
    "opaque" ->
          (SubExp -> InternaliseM SubExp)
    -> [SubExp] -> InternaliseM [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM ([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
          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'
          forM (zip ts e') $ \(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
          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'
          forM (zip ts e') $ \(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
      e1' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "assert_cond" Exp
    e1
      c <- assert "assert_c" e1' $ errorMsg [ErrorString $ "Assertion is false: " <> check]
      -- Make sure there are some bindings to certify.
      certifying c $ mapM rebind =<< internaliseExp desc e2
      where
        rebind :: SubExp -> m SubExp
    rebind SubExp
    v = do
          v' <- [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "assert_res"
          letBindNames [v'] $ I.BasicOp $ I.SubExp v
          pure $ I.Var 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
      (ts, 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
      es' <- concat <$> mapM (internaliseExp "payload") es

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

      case lookupWithIndex c constr_map of
        Just (Int
    i, [Int]
    js) ->
          (IntType -> Integer -> SubExp
    intConst IntType
    Int8 (Int -> Integer
    forall a. Integral a => a -> Integer
    toInteger Int
    i) SubExp -> [SubExp] -> [SubExp]
    forall a. a -> [a] -> [a]
    :) ([SubExp] -> [SubExp])
    -> InternaliseM [SubExp] -> InternaliseM [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Int
    -> [TypeBase Shape 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] -> [SubExp]
    forall a. a -> [a] -> [a]
    :) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> t -> [TypeBase Shape 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
              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
              (blank :) <$> clauses (j + 1) ts js_to_es
        clauses t
    _ [] [(t, SubExp)]
    _ =
          [SubExp] -> f [SubExp]
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure []
    internaliseExp [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
      exists <- InternaliseM (Scope SOACS)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
      case argdim of
        Just VName
    d | VName
    d VName -> Scope SOACS -> Bool
    forall k a. Ord k => k -> Map k a -> Bool
    `M.member` Scope SOACS
    exists -> [SubExp] -> InternaliseM [SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [VName -> SubExp
    I.Var VName
    d]
        Maybe VName
    _ -> do
          arg' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc Exp
    arg
          case (arg', argdim) of
            ([SubExp
    se], Just VName
    d) -> do
              [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    d] (Exp (Rep InternaliseM) -> InternaliseM ())
    -> Exp (Rep InternaliseM) -> InternaliseM ()
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep InternaliseM)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep InternaliseM))
    -> BasicOp -> Exp (Rep InternaliseM)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp SubExp
    se
            ([SubExp], Maybe VName)
    _ -> () -> InternaliseM ()
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
          pure arg'

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

    generateCond ::
      E.Pat StructType ->
      [I.SubExp] ->
      InternaliseM ([Maybe I.PrimValue], [I.SubExp])
    generateCond :: PatBase Info VName StructType
    -> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
    generateCond PatBase Info VName StructType
    orig_p [SubExp]
    orig_ses = do
      (cmps, pertinent, _) <- PatBase Info VName StructType
    -> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp], [SubExp])
    forall {vn} {a}.
    IsName vn =>
    PatBase Info vn StructType
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    compares PatBase Info VName StructType
    orig_p [SubExp]
    orig_ses
      pure (cmps, pertinent)
      where
        compares :: PatBase Info vn StructType
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    compares (E.PatLit PatLit
    l (Info StructType
    t) SrcLoc
    _) (a
    se : [a]
    ses) =
          ([Maybe PrimValue], [a], [a])
    -> InternaliseM ([Maybe PrimValue], [a], [a])
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([PrimValue -> Maybe PrimValue
    forall a. a -> Maybe a
    Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
    forall a b. (a -> b) -> a -> b
    $ PatLit -> StructType -> PrimValue
    internalisePatLit PatLit
    l StructType
    t], [a
    se], [a]
    ses)
        compares (E.PatConstr Name
    c (Info (E.Scalar (E.Sum Map Name [StructType]
    fs))) [PatBase Info vn StructType]
    pats SrcLoc
    _) (a
    _ : [a]
    ses) = do
          (payload_ts, 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 lookupWithIndex c m of
            Just (Int
    tag, [Int]
    payload_is) -> do
              let ([a]
    payload_ses, [a]
    ses') = Int -> [a] -> ([a], [a])
    forall a. Int -> [a] -> ([a], [a])
    splitAt ([TypeBase ExtShape Uniqueness] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TypeBase ExtShape Uniqueness]
    payload_ts) [a]
    ses
              (cmps, pertinent, _) <-
                [PatBase Info vn StructType]
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    comparesMany [PatBase Info vn StructType]
    pats ([a] -> InternaliseM ([Maybe PrimValue], [a], [a]))
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    forall a b. (a -> b) -> a -> b
    $ (Int -> a) -> [Int] -> [a]
    forall a b. (a -> b) -> [a] -> [b]
    map ([a]
    payload_ses [a] -> Int -> a
    forall a. HasCallStack => [a] -> Int -> a
    !!) [Int]
    payload_is
              let missingCmps Int
    i a
    _ =
                    case Int
    i Int -> [Int] -> Maybe Int
    forall a. Eq a => a -> [a] -> Maybe Int
    `elemIndex` [Int]
    payload_is of
                      Just Int
    j -> [Maybe PrimValue]
    cmps [Maybe PrimValue] -> Int -> Maybe PrimValue
    forall a. HasCallStack => [a] -> Int -> a
    !! Int
    j
                      Maybe Int
    Nothing -> Maybe PrimValue
    forall a. Maybe a
    Nothing
              pure
                ( Just (I.IntValue $ intValue Int8 $ toInteger tag)
                    : zipWith missingCmps [0 ..] payload_ses,
                  pertinent,
                  ses'
                )
            Maybe (Int, [Int])
    Nothing ->
              [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "generateCond: missing constructor"
        compares (E.PatConstr Name
    _ (Info StructType
    t) [PatBase Info vn StructType]
    _ SrcLoc
    _) [a]
    _ =
          [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> InternaliseM ([Maybe PrimValue], [a], [a]))
    -> [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "generateCond: PatConstr has nonsensical type: " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ StructType -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString StructType
    t
        compares (E.Id vn
    _ Info StructType
    t SrcLoc
    loc) [a]
    ses =
          PatBase Info vn StructType
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    compares (Info StructType -> SrcLoc -> PatBase Info vn StructType
    forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
    E.Wildcard Info StructType
    t SrcLoc
    loc) [a]
    ses
        compares (E.Wildcard (Info StructType
    t) SrcLoc
    _) [a]
    ses = do
          let ([a]
    id_ses, [a]
    rest_ses) = Int -> [a] -> ([a], [a])
    forall a. Int -> [a] -> ([a], [a])
    splitAt (StructType -> Int
    forall 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
          (cmps1, pertinent1, ses') <- PatBase Info vn StructType
    -> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
    compares PatBase Info vn StructType
    pat [a]
    ses
          (cmps2, pertinent2, ses'') <- comparesMany pats ses'
          pure
            ( cmps1 <> cmps2,
              pertinent1 <> pertinent2,
              ses''
            )

    internalisePat ::
      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
      ses <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    desc' Exp
    e
      internalisePat' sizes p ses 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
      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
      stmPat (toParam E.Observe <$> p) ses_ts $ \[VName]
    pat_names -> do
        AppRes -> [SubExp] -> InternaliseM ()
    bindExtSizes (StructType -> [VName] -> AppRes
    AppRes (PatBase Info VName StructType -> StructType
    forall d 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
      (idxs', oks, parts) <- [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
    -> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 ([(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
     -> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]]))
    -> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
    -> InternaliseM
         ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp
     -> DimIndex
     -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]))
    -> [SubExp]
    -> SliceBase Info VName
    -> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m [c]
    zipWithM SubExp
    -> DimIndex
    -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
    internaliseDimIndex [SubExp]
    dims SliceBase Info VName
    idxs
      ok <- letSubExp "index_ok" =<< eAll oks
      let msg =
            [ErrorMsgPart SubExp] -> ErrorMsg SubExp
    forall a. [ErrorMsgPart a] -> ErrorMsg a
    errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
    -> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
    forall a b. (a -> b) -> a -> b
    $
              [ErrorMsgPart SubExp
    "Index ["]
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [ErrorMsgPart SubExp]
    -> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [[a]] -> [a]
    intercalate [ErrorMsgPart SubExp
    ", "] [[ErrorMsgPart SubExp]]
    parts
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [ErrorMsgPart SubExp
    "] out of bounds for array of shape ["]
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ ErrorMsgPart SubExp
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. a -> [a] -> [a]
    intersperse ErrorMsgPart SubExp
    "][" ((SubExp -> ErrorMsgPart SubExp)
    -> [SubExp] -> [ErrorMsgPart SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (PrimType -> SubExp -> ErrorMsgPart SubExp
    forall a. PrimType -> a -> ErrorMsgPart a
    ErrorVal PrimType
    int64) ([SubExp] -> [ErrorMsgPart SubExp])
    -> [SubExp] -> [ErrorMsgPart SubExp]
    forall a b. (a -> b) -> a -> b
    $ Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take (SliceBase Info VName -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length SliceBase Info VName
    idxs) [SubExp]
    dims)
                [ErrorMsgPart SubExp]
    -> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [ErrorMsgPart SubExp
    "]."]
      c <- assert "index_certs" ok msg
      pure (idxs', c)

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

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

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

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

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

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

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

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

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

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

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

    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
      arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars ([Char]
    what [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_arr") Exp
    arr
      nes <- internaliseExp (what ++ "_ne") ne
      nes' <- forM (zip nes arrs) $ \(SubExp
    ne', VName
    arr') -> do
        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'
        ensureShape
          "Row shape of input array does not match shape of neutral element"
          rowtype
          (what ++ "_ne_right_shape")
          ne'
      nests <- mapM I.subExpType nes'
      arrts <- mapM lookupType arrs
      lam' <- internaliseFoldLambda internaliseLambda lam nests arrts
      w <- arraysSize 0 <$> mapM lookupType arrs
      letValExp' desc . I.Op =<< f w lam' nes' 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
      rf' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "hist_rf" Exp
    rf
      ne' <- internaliseExp "hist_ne" ne
      hist' <- internaliseExpToVars "hist_hist" hist
      buckets' <- internaliseExpToVars "hist_buckets" buckets
      img' <- internaliseExpToVars "hist_img" img

      -- reshape neutral element to have same size as the destination array
      ne_shp <- forM (zip ne' hist') $ \(SubExp
    n, VName
    h) -> do
        rowtype <- Int -> TypeBase Shape 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
        ensureShape
          "Row shape of destination array does not match shape of neutral element"
          rowtype
          "hist_ne_right_shape"
          n
      ne_ts <- mapM I.subExpType ne_shp
      his_ts <- mapM (fmap (I.stripArray (dim - 1)) . lookupType) hist'
      op' <- internaliseFoldLambda internaliseLambda op ne_ts his_ts

      -- reshape return type of bucket function to have same size as neutral element
      -- (modulo the index)
      bucket_params <- replicateM dim (newParam "bucket_p" $ I.Prim int64)
      img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img'
      let params = [Param (TypeBase Shape 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 = 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 = 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
      lam' <-
        mkLambda params $
          ensureResultShape
            "Row shape of value array does not match row shape of hist target"
            rettype
            =<< bodyBind body

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

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

    internaliseStreamAcc ::
      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
      dest' <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "scatter_dest" Exp
    dest
      bs' <- internaliseExpToVars "scatter_input" bs

      acc_cert_v <- newVName "acc_cert"
      dest_ts <- mapM lookupType dest'
      let dest_w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    dest_ts
          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
      acc_p <- newParam "acc_p" acc_t
      withacc_lam <- mkLambda [Param mempty acc_cert_v (I.Prim I.Unit), acc_p] $ do
        bs_ts <- mapM lookupType bs'
        lam' <- internaliseLambdaCoerce lam $ map rowType $ paramType acc_p : bs_ts
        let w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    bs_ts
        fmap subExpsRes . letValExp' "acc_res" $
          I.Op $
            I.Screma w (I.paramName acc_p : bs') (I.mapSOAC lam')

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

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

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

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

    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
      vs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    s Exp
    e
      mapM (letSubExp s . I.BasicOp <=< op) vs

    certifyingNonzero ::
      IntType ->
      SubExp ->
      InternaliseM a ->
      InternaliseM a
    certifyingNonzero :: forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
    certifyingNonzero IntType
    t SubExp
    x InternaliseM a
    m = do
      zero <-
        [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)
      nonzero <- letSubExp "nonzero" $ I.BasicOp $ UnOp (I.Neg I.Bool) zero
      c <- assert "nonzero_cert" nonzero "division by zero"
      certifying c m

    certifyingNonnegative ::
      IntType ->
      SubExp ->
      InternaliseM a ->
      InternaliseM a
    certifyingNonnegative :: forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
    certifyingNonnegative IntType
    t SubExp
    x InternaliseM a
    m = do
      nonnegative <-
        [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
      c <- assert "nonzero_cert" nonnegative "negative exponent"
      certifying c 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
      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
      fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp (I.Neg I.Bool) 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' <- [Char] -> Exp -> InternaliseM (Body SOACS)
    internaliseBody [Char]
    "lam" Exp
    body
        rettype' <- internaliseLambdaReturnType rettype =<< bodyExtType body'
        pure (params', body', 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
      (params, body, rettype) <- InternaliseLambda
    internaliseLambda Exp
    lam [TypeBase Shape NoUniqueness]
    argtypes
      mkLambda params $
        ensureResultShape
          (ErrorMsg [ErrorString "unexpected lambda result size"])
          rettype
          =<< bodyBind 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
              rs <- (SubExp -> SubExp -> InternaliseM SubExp)
    -> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m [c]
    zipWithM SubExp -> SubExp -> InternaliseM SubExp
    doComparison [SubExp]
    xe' [SubExp]
    ye'
              cmp_f =<< letSubExp "eq" =<< eAll rs
          where
            isEqlOp :: [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
              x_t <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
    forall t (m :: * -> *).
    HasScope t m =>
    SubExp -> m (TypeBase Shape NoUniqueness)
    I.subExpType SubExp
    x
              y_t <- I.subExpType y
              case 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
                  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
                  shapes_match <- letSubExp "shapes_match" =<< eAll dims_match
                  let 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.
                        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
                        x' <- letExp "x" $ I.BasicOp $ I.SubExp x
                        y' <- letExp "x" $ I.BasicOp $ I.SubExp y
                        x_flat <-
                          letExp "x_flat" . I.BasicOp $
                            I.Reshape x' (reshapeAll (I.arrayShape x_t) (I.Shape [x_num_elems]))
                        y_flat <-
                          letExp "y_flat" . I.BasicOp $
                            I.Reshape y' (reshapeAll (I.arrayShape x_t) (I.Shape [x_num_elems]))

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

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

                  letSubExp "arrays_equal"
                    =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])
        handle [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] -> [Char] -> Bool
    forall a. Eq a => a -> a -> Bool
    ==) ([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
              x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop 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
              x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              y' <- internaliseExp1 "y" y
              fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'
          | Just CmpOp
    cmp <- (CmpOp -> Bool) -> [CmpOp] -> Maybe CmpOp
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (([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
              x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              y' <- internaliseExp1 "y" y
              fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' 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
              x' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "x" Exp
    x
              fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv 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
          arr' <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "map_arr" Exp
    arr
          arr_ts <- mapM lookupType arr'
          lam' <- internaliseLambdaCoerce lam $ map rowType arr_ts
          let w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    arr_ts
          letTupExp' desc $ I.Op $ I.Screma w arr' (I.mapSOAC lam')
        handleSOACs [Exp
    k, Exp
    lam, Exp
    arr] a
    "partition" = do
          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
          Just $ \[Char]
    _desc -> do
            arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "partition_input" Exp
    arr
            lam' <- internalisePartitionLambda internaliseLambda k' lam $ map I.Var arrs
            uncurry (++) <$> partitionWithSOACS (fromIntegral k') lam' 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
          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
          i' <- internaliseExp1 "acc_i" i
          vs <- internaliseExp "acc_v" v
          fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc Safe acc' [i'] 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
              x' <- [Char] -> Exp -> InternaliseM [SubExp]
    internaliseExp [Char]
    "ad_x" Exp
    x
              v' <- internaliseExp "ad_v" v
              lam <- internaliseLambdaCoerce f =<< mapM subExpType x'
              fmap (map I.Var) . letTupExp desc . Op $
                case 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
          arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "unflatten_arr" Exp
    arr
          n' <- internaliseExp1 "n" n
          m' <- internaliseExp1 "m" m
          -- Each dimension must be nonnegative, and the unflattened
          -- dimension needs to have the same number of elements as the
          -- original dimension.
          old_dim <- I.arraysSize 0 <$> mapM lookupType arrs
          dim_ok <-
            letSubExp "dim_ok" <=< toExp $
              pe64 old_dim .==. pe64 n'
                * pe64 m'
                  .&&. pe64 n'
                  .>=. 0
                  .&&. pe64 m'
                  .>=. 0
          dim_ok_cert <-
            assert "dim_ok_cert" dim_ok $
              ErrorMsg
                [ "Cannot unflatten array of shape [",
                  ErrorVal int64 old_dim,
                  "] to array of shape [",
                  ErrorVal int64 n',
                  "][",
                  ErrorVal int64 m',
                  "]"
                ]
          certifying dim_ok_cert $
            forM arrs $ \VName
    arr' -> do
              arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    arr'
              letSubExp desc . I.BasicOp $
                I.Reshape arr' $
                  reshapeAll (I.arrayShape arr_t) $
                    reshapeOuter (I.Shape [n', m']) 1 $
                      I.arrayShape arr_t
        handleRest [Exp
    arr] a
    "manifest" = ([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
          arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "flatten_arr" Exp
    arr
          forM arrs $ \VName
    arr' -> do
            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 r == 0
              then pure $ I.Var arr'
              else letSubExp desc $ I.BasicOp $ I.Manifest arr' [0 .. r - 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
          arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "flatten_arr" Exp
    arr
          forM arrs $ \VName
    arr' -> do
            arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    arr'
            let n = Int -> TypeBase Shape NoUniqueness -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    arraySize Int
    0 TypeBase Shape NoUniqueness
    arr_t
                m = Int -> TypeBase Shape NoUniqueness -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    arraySize Int
    1 TypeBase Shape NoUniqueness
    arr_t
            k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowUndef) n m
            letSubExp desc . I.BasicOp $
              I.Reshape arr' $
                reshapeAll (I.arrayShape arr_t) $
                  reshapeOuter (I.Shape [k]) 2 $
                    I.arrayShape arr_t
        handleRest [Exp
    x, Exp
    y] a
    "concat" = ([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
          xs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "concat_x" Exp
    x
          ys <- internaliseExpToVars "concat_y" y
          outer_size <- arraysSize 0 <$> mapM lookupType xs
          let 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
          ressize <-
            foldM sumdims outer_size
              =<< mapM (fmap (arraysSize 0) . mapM lookupType) [ys]

          let conc VName
    xarr VName
    yarr =
                BasicOp -> Exp SOACS
    forall rep. BasicOp -> Exp rep
    I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $ Int -> NonEmpty VName -> SubExp -> BasicOp
    I.Concat Int
    0 (VName
    xarr VName -> [VName] -> NonEmpty VName
    forall a. a -> [a] -> NonEmpty a
    :| [VName
    yarr]) SubExp
    ressize
          mapM (letSubExp desc) $ zipWith conc xs ys
        handleRest [Exp
    e] a
    "transpose" = ([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
            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
            pure $ I.Rearrange v ([1, 0] ++ [2 .. r - 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
          e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "trunc_arg" Exp
    e
          case E.typeOf 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
          e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "trunc_arg" Exp
    e
          case E.typeOf 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
          si' <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "write_arg_i" Exp
    si
          svs <- internaliseExpToVars "write_arg_v" v
          sas <- internaliseExpToVars "write_arg_a" a

          si_w <- I.arraysSize 0 <$> mapM lookupType si'
          sv_ts <- mapM lookupType svs

          svs' <- forM (zip svs sv_ts) $ \(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.
            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
            c <-
              assert
                "write_cert"
                cmp
                "length of index and value array does not match"
            certifying c $
              letExp (baseString sv ++ "_write_sv") $
                shapeCoerce (I.shapeDims (reshapeOuter (I.Shape [si_w]) 1 sv_shape)) sv

          indexType <- fmap rowType <$> mapM lookupType si'
          indexName <- mapM (\TypeBase Shape NoUniqueness
    _ -> [Char] -> InternaliseM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "write_index") indexType
          valueNames <- replicateM (length sv_ts) $ newVName "write_value"

          sa_ts <- mapM lookupType sas
          let 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]
    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]
    indexName [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    valueNames
              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 <- localScope (scopeOfLParams bodyParams) . buildBody_ $ do
            let 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
            results <- forM outs $ \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
            ensureResultShape
              "scatter value has wrong size"
              bodyTypes
              (subExpsRes results)

          let 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]
    si' [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    svs'

          let 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] -> [(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
          letTupExp' desc $ I.Op $ I.Scatter si_w sivs spec 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
      arrs <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "arr" Exp
    arr
      offset' <- internaliseExp1 "offset" offset
      old_dim <- I.arraysSize 0 <$> mapM lookupType arrs
      offset_inbounds_down <- letSubExp "offset_inbounds_down" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) offset'
      offset_inbounds_up <- letSubExp "offset_inbounds_up" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) offset' old_dim
      slices' <-
        mapM
          ( \(Exp
    n, Exp
    s) -> do
              n' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "n" Exp
    n
              s' <- internaliseExp1 "s" s
              pure (n', s')
          )
          slices
      (min_bound, max_bound) <-
        foldM
          ( \(SubExp
    lower, SubExp
    upper) (SubExp
    n, SubExp
    s) -> do
              n_m1 <- [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)
              spn <- letSubExp "span" $ I.BasicOp $ I.BinOp (I.Mul Int64 I.OverflowUndef) n_m1 s

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

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

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

      all_bounds <-
        foldM
          (\SubExp
    x SubExp
    y -> [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)
          offset_inbounds_down
          [offset_inbounds_up, min_in_bounds, max_in_bounds]

      c <-
        assert "bounds_cert" all_bounds $
          ErrorMsg [ErrorString $ "Flat slice out of bounds: " <> prettyText old_dim <> " and " <> prettyText slices']
      let slice = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
    forall d. d -> [FlatDimIndex d] -> FlatSlice d
    I.FlatSlice SubExp
    offset' ([FlatDimIndex SubExp] -> FlatSlice SubExp)
    -> [FlatDimIndex SubExp] -> FlatSlice SubExp
    forall a b. (a -> b) -> a -> b
    $ ((SubExp, SubExp) -> FlatDimIndex SubExp)
    -> [(SubExp, SubExp)] -> [FlatDimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map ((SubExp -> SubExp -> FlatDimIndex SubExp)
    -> (SubExp, SubExp) -> FlatDimIndex SubExp
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry SubExp -> SubExp -> FlatDimIndex SubExp
    forall d. d -> d -> FlatDimIndex d
    FlatDimIndex) [(SubExp, SubExp)]
    slices'
      certifying c $
        forM arrs $ \VName
    arr' ->
          [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
      arrs1 <- [Char] -> Exp -> InternaliseM [VName]
    internaliseExpToVars [Char]
    "arr" Exp
    arr1
      offset' <- internaliseExp1 "offset" offset
      old_dim <- I.arraysSize 0 <$> mapM lookupType arrs1
      offset_inbounds_down <- letSubExp "offset_inbounds_down" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) offset'
      offset_inbounds_up <- letSubExp "offset_inbounds_up" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) offset' old_dim
      arrs2 <- internaliseExpToVars "arr" arr2
      ts <- mapM lookupType arrs2
      slices' <-
        mapM
          ( \(Exp
    s, Int
    i) -> do
              s' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "s" Exp
    s
              let n = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    i [TypeBase Shape NoUniqueness]
    ts
              pure (n, s')
          )
          $ zip slices [0 ..]
      (min_bound, max_bound) <-
        foldM
          ( \(SubExp
    lower, SubExp
    upper) (SubExp
    n, SubExp
    s) -> do
              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)
              spn <- letSubExp "span" $ I.BasicOp $ I.BinOp (I.Mul Int64 I.OverflowUndef) n_m1 s

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

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

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

      all_bounds <-
        foldM
          (\SubExp
    x SubExp
    y -> [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)
          offset_inbounds_down
          [offset_inbounds_up, min_in_bounds, max_in_bounds]

      c <-
        assert "bounds_cert" all_bounds $
          ErrorMsg [ErrorString $ "Flat slice out of bounds: " <> prettyText old_dim <> " and " <> prettyText slices']
      let slice = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
    forall d. d -> [FlatDimIndex d] -> FlatSlice d
    I.FlatSlice SubExp
    offset' ([FlatDimIndex SubExp] -> FlatSlice SubExp)
    -> [FlatDimIndex SubExp] -> FlatSlice SubExp
    forall a b. (a -> b) -> a -> b
    $ ((SubExp, SubExp) -> FlatDimIndex SubExp)
    -> [(SubExp, SubExp)] -> [FlatDimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map ((SubExp -> SubExp -> FlatDimIndex SubExp)
    -> (SubExp, SubExp) -> FlatDimIndex SubExp
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry SubExp -> SubExp -> FlatDimIndex SubExp
    forall d. d -> d -> FlatDimIndex d
    FlatDimIndex) [(SubExp, SubExp)]
    slices'
      certifying c $
        forM (zip arrs1 arrs2) $ \(VName
    arr1', VName
    arr2') ->
          [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
      (shapes, value_paramts, fun_params, rettype_fun) <- VName -> InternaliseM FunInfo
    lookupFunction VName
    fname
      argts <- mapM subExpType args

      shapeargs <- argShapes shapes fun_params argts
      let diets =
            Int -> Diet -> [Diet]
    forall a. Int -> a -> [a]
    replicate ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    shapeargs) Diet
    I.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
      args' <-
        ensureArgShapes
          "function arguments of wrong shape"
          (map I.paramName fun_params)
          (map I.paramType fun_params)
          (shapeargs ++ args)
      argts' <- mapM subExpType args'
      case rettype_fun $ zip args' argts' of
        Maybe [(TypeBase ExtShape 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 <- InternaliseM Safety
    askSafety
          attrs <- asks envAttrs
          attributing attrs . letValExp' desc $
            I.Apply (internaliseFunName fname) (zip args' diets) ts safety

    -- Bind existential names defined by an expression, based on the
    -- concrete values that expression evaluated to.  This most
    -- importantly should be done after function calls, but also
    -- everything else that can produce existentials in the source
    -- language.
    bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
    bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
    bindExtSizes (AppRes StructType
    ret [VName]
    retext) [SubExp]
    ses = do
      let ts :: [TypeBase ExtShape 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
      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
    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' (I.Free (I.Var VName
    v)) SubExp
    se
            | VName
    v VName -> [VName] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` [VName]
    retext = VName -> SubExp -> Map VName SubExp
    forall k a. k -> a -> Map k a
    M.singleton VName
    v SubExp
    se
          combine' Ext SubExp
    _ SubExp
    _ = Map VName SubExp
    forall a. Monoid a => a
    mempty

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

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

    -- 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
      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 = Int -> [TypeBase Shape NoUniqueness] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [TypeBase Shape NoUniqueness]
    arr_ts
      classes_and_increments <- letTupExp "increments" $ I.Op $ I.Screma w arrs (mapSOAC lam)
      (classes, increments) <- case 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"

      add_lam_x_params <-
        replicateM k $ newParam "x" (I.Prim int64)
      add_lam_y_params <-
        replicateM k $ newParam "y" (I.Prim int64)
      add_lam_body <- runBodyBuilder $
        localScope (scopeOfLParams $ add_lam_x_params ++ add_lam_y_params) $
          fmap subExpsRes $
            forM (zip add_lam_x_params add_lam_y_params) $ \(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 =
            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 = 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

      scan <- I.scanSOAC [I.Scan add_lam nes]
      all_offsets <- letTupExp "offsets" $ I.Op $ I.Screma w increments 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.
      last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int64 OverflowUndef) w $ constant (1 :: Int64)
      let 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 = [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)
      is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int64) w $ constant (0 :: Int64)
      sizes <-
        letTupExp "partition_size" =<< eIf (eSubExp is_empty) empty_body 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.
      blanks <- forM arr_ts $ \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.
      write_lam <- do
        c_param <- newParam "c" (I.Prim int64)
        offset_params <- replicateM k $ newParam "offset" (I.Prim int64)
        value_params <- mapM (newParam "v" . I.rowType) arr_ts
        (offset, offset_stms) <-
          collectStms $
            mkOffsetLambdaBody
              (map I.Var sizes)
              (I.Var $ I.paramName c_param)
              0
              offset_params
        pure
          I.Lambda
            { I.lambdaParams = c_param : offset_params ++ value_params,
              I.lambdaReturnType =
                replicate (length arr_ts) (I.Prim int64)
                  ++ map I.rowType arr_ts,
              I.lambdaBody =
                mkBody offset_stms $
                  replicate (length arr_ts) (subExpRes offset)
                    ++ I.varsRes (map I.paramName value_params)
            }
      let 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
      results <-
        letTupExp "partition_res" . I.Op $
          I.Scatter w (classes : all_offsets ++ arrs) spec write_lam
      sizes' <-
        letSubExp "partition_sizes" $
          I.BasicOp $
            I.ArrayLit (map I.Var sizes) $
              I.Prim int64
      pure (map I.Var results, [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
          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
          next_one <- mkOffsetLambdaBody sizes c (i + 1) ps
          this_one <-
            letSubExp "this_offset"
              =<< foldBinOp
                (Add Int64 OverflowUndef)
                (constant (-1 :: Int64))
                (I.Var (I.paramName p) : take i sizes)
          letSubExp "total_res"
            =<< eIf
              (eSubExp is_this_one)
              (resultBodyM [this_one])
              (resultBodyM [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
          e' <- [Char] -> Exp -> InternaliseM SubExp
    internaliseExp1 [Char]
    "size" Exp
    e
          pure ["[", ErrorVal int64 e', "]"]

    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
      args' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
    -> InternaliseM [[ErrorMsgPart SubExp]]
    -> InternaliseM [ErrorMsgPart SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp])
    -> [TypeArg Exp] -> InternaliseM [[ErrorMsgPart SubExp]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
    onArg [TypeArg Exp]
    args
      pure $ intersperse " " $ ErrorString (prettyText v) : args'
      where
        onArg :: TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
    onArg (TypeArgDim Exp
    d) = Exp -> InternaliseM [ErrorMsgPart SubExp]
    sizeExpForError Exp
    d
        onArg (TypeArgType StructType
    t) = StructType -> InternaliseM [ErrorMsgPart SubExp]
    forall 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
          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
          pure $ ["("] ++ intercalate [", "] ts' ++ [")"]
      | Bool
    otherwise = do
          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
          pure $ ["{"] ++ intercalate [", "] fs' ++ ["}"]
      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] -> [ErrorMsgPart SubExp]
    forall a. a -> [a] -> [a]
    :) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
    -> InternaliseM [ErrorMsgPart SubExp]
    -> InternaliseM [ErrorMsgPart SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> TypeBase Exp 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
      shape' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
    forall a. Monoid a => [a] -> a
    mconcat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
    -> InternaliseM [[ErrorMsgPart SubExp]]
    -> InternaliseM [ErrorMsgPart SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> InternaliseM [ErrorMsgPart SubExp])
    -> [Exp] -> InternaliseM [[ErrorMsgPart SubExp]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM Exp -> InternaliseM [ErrorMsgPart SubExp]
    sizeExpForError (Shape Exp -> [Exp]
    forall dim. Shape dim -> [dim]
    E.shapeDims Shape Exp
    shape)
      et' <- typeExpForError $ Scalar et
      pure $ shape' ++ et'
    typeExpForError (E.Scalar (E.Sum Map Name [TypeBase Exp u]
    cs)) = do
      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
      pure $ intercalate [" | "] cs'
      where
        onConstructor :: (a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
    onConstructor (a
    c, [TypeBase Exp u]
    ts) = do
          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
          pure $ ErrorString ("#" <> prettyText c <> " ") : intercalate [" "] ts'
    typeExpForError (E.Scalar Arrow {}) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [ErrorMsgPart SubExp
    "#<fun>"]

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

    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
    "]"