    {-# LANGUAGE TypeFamilies #-}

    -- | = Constructing Futhark ASTs
    --
    -- This module re-exports and defines a bunch of building blocks for
    -- constructing fragments of Futhark ASTs.  More importantly, it also
    -- contains a basic introduction on how to use them.
    --
    -- The "Futhark.IR.Syntax" module contains the core
    -- AST definition.  One important invariant is that all bound names in
    -- a Futhark program must be /globally/ unique.  In principle, you
    -- could use the facilities from "Futhark.MonadFreshNames" (or your
    -- own bespoke source of unique names) to manually construct
    -- expressions, statements, and entire ASTs.  In practice, this would
    -- be very tedious.  Instead, we have defined a collection of building
    -- blocks (centered around the 'MonadBuilder' type class) that permits
    -- a more abstract way of generating code.
    --
    -- Constructing ASTs with these building blocks requires you to ensure
    -- that all free variables are in scope.  See
    -- "Futhark.IR.Prop.Scope".
    --
    -- == 'MonadBuilder'
    --
    -- A monad that implements 'MonadBuilder' tracks the statements added
    -- so far, the current names in scope, and allows you to add
    -- additional statements with 'addStm'.  Any monad that implements
    -- 'MonadBuilder' also implements the t'Rep' type family, which
    -- indicates which rep it works with.  Inside a 'MonadBuilder' we can
    -- use 'collectStms' to gather up the 'Stms' added with 'addStm' in
    -- some nested computation.
    --
    -- The 'BuilderT' monad (and its convenient 'Builder' version) provides
    -- the simplest implementation of 'MonadBuilder'.
    --
    -- == Higher-level building blocks
    --
    -- On top of the raw facilities provided by 'MonadBuilder', we have
    -- more convenient facilities.  For example, 'letSubExp' lets us
    -- conveniently create a 'Stm' for an 'Exp' that produces a /single/
    -- value, and returns the (fresh) name for the resulting variable:
    --
    -- @
    -- z <- letExp "z" $ BasicOp $ BinOp (Add Int32) (Var x) (Var y)
    -- @
    --
    -- == Monadic expression builders
    --
    -- This module also contains "monadic expression" functions that let
    -- us build nested expressions in a "direct" style, rather than using
    -- 'letExp' and friends to bind every sub-part first.  See functions
    -- such as 'eIf' and 'eBody' for example.  See also
    -- "Futhark.Analysis.PrimExp" and the 'ToExp' type class.
    --
    -- == Examples
    --
    -- The "Futhark.Transform.FirstOrderTransform" module is a
    -- (relatively) simple example of how to use these components.  As are
    -- some of the high-level building blocks in this very module.
    module Futhark.Construct
      ( -- * Basic building blocks
        module Futhark.Builder,
        letSubExp,
        letExp,
        letTupExp,
        letTupExp',
        letInPlace,

        -- * Monadic expression builders
        eSubExp,
        eParam,
        eMatch',
        eMatch,
        eIf,
        eIf',
        eBinOp,
        eUnOp,
        eCmpOp,
        eConvOp,
        eSignum,
        eCopy,
        eBody,
        eLambda,
        eBlank,
        eAll,
        eAny,
        eDimInBounds,
        eOutOfBounds,
        eIndex,
        eLast,

        -- * Other building blocks
        asIntZ,
        asIntS,
        resultBody,
        resultBodyM,
        insertStmsM,
        buildBody,
        buildBody_,
        mapResult,
        foldBinOp,
        binOpLambda,
        cmpOpLambda,
        mkLambda,
        sliceDim,
        fullSlice,
        fullSliceNum,
        isFullSlice,
        sliceAt,

        -- * Result types
        instantiateShapes,
        instantiateShapes',
        removeExistentials,

        -- * Convenience
        simpleMkLetNames,
        ToExp (..),
        toSubExp,
      )
    where

    import Control.Monad
    import Control.Monad.Identity
    import Control.Monad.State
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Futhark.Builder
    import Futhark.IR
    import Futhark.Util (maybeNth)

    -- | @letSubExp desc e@ binds the expression @e@, which must produce a
    -- single value.  Returns a t'SubExp' corresponding to the resulting
    -- value.  For expressions that produce multiple values, see
    -- 'letTupExp'.
    letSubExp ::
      (MonadBuilder m) =>
      String ->
      Exp (Rep m) ->
      m SubExp
    letSubExp :: forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    _ (BasicOp (SubExp SubExp
    se)) = SubExp -> m SubExp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    se
    letSubExp String
    desc Exp (Rep m)
    e = VName -> SubExp
    Var (VName -> SubExp) -> m VName -> m SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> Exp (Rep m) -> m VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp String
    desc Exp (Rep m)
    e

    -- | Like 'letSubExp', but returns a name rather than a t'SubExp'.
    letExp ::
      (MonadBuilder m) =>
      String ->
      Exp (Rep m) ->
      m VName
    letExp :: forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp String
    _ (BasicOp (SubExp (Var VName
    v))) =
      VName -> m VName
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    v
    letExp String
    desc Exp (Rep m)
    e = do
      n <- [ExtType] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([ExtType] -> Int) -> m [ExtType] -> m Int
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp (Rep m) -> m [ExtType]
    forall rep (m :: * -> *).
    (HasScope rep m, TypedOp (OpC rep)) =>
    Exp rep -> m [ExtType]
    expExtType Exp (Rep m)
    e
      vs <- replicateM n $ newVName desc
      letBindNames vs e
      case vs of
        [VName
    v] -> VName -> m VName
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    v
        [VName]
    _ -> String -> m VName
    forall a. HasCallStack => String -> a
    error (String -> m VName) -> String -> m VName
    forall a b. (a -> b) -> a -> b
    $ String
    "letExp: tuple-typed expression given:\n" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Exp (Rep m) -> String
    forall a. Pretty a => a -> String
    prettyString Exp (Rep m)
    e

    -- | Like 'letExp', but the 'VName' and 'Slice' denote an array that
    -- is 'Update'd with the result of the expression.  The name of the
    -- updated array is returned.
    letInPlace ::
      (MonadBuilder m) =>
      String ->
      VName ->
      Slice SubExp ->
      Exp (Rep m) ->
      m VName
    letInPlace :: forall (m :: * -> *).
    MonadBuilder m =>
    String -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
    letInPlace String
    desc VName
    src Slice SubExp
    slice Exp (Rep m)
    e = do
      tmp <- String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp (String
    desc String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    "_tmp") Exp (Rep m)
    e
      letExp desc $ BasicOp $ Update Unsafe src slice tmp

    -- | Like 'letExp', but the expression may return multiple values.
    letTupExp ::
      (MonadBuilder m) =>
      String ->
      Exp (Rep m) ->
      m [VName]
    letTupExp :: forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m [VName]
    letTupExp String
    _ (BasicOp (SubExp (Var VName
    v))) =
      [VName] -> m [VName]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [VName
    v]
    letTupExp String
    name Exp (Rep m)
    e = do
      e_t <- Exp (Rep m) -> m [ExtType]
    forall rep (m :: * -> *).
    (HasScope rep m, TypedOp (OpC rep)) =>
    Exp rep -> m [ExtType]
    expExtType Exp (Rep m)
    e
      names <- replicateM (length e_t) $ newVName name
      letBindNames names e
      pure names

    -- | Like 'letTupExp', but returns t'SubExp's instead of 'VName's.
    letTupExp' ::
      (MonadBuilder m) =>
      String ->
      Exp (Rep m) ->
      m [SubExp]
    letTupExp' :: forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m [SubExp]
    letTupExp' String
    _ (BasicOp (SubExp SubExp
    se)) = [SubExp] -> m [SubExp]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp
    se]
    letTupExp' String
    name Exp (Rep m)
    ses = (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var ([VName] -> [SubExp]) -> m [VName] -> m [SubExp]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> Exp (Rep m) -> m [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m [VName]
    letTupExp String
    name Exp (Rep m)
    ses

    -- | Turn a subexpression into a monad expression.  Does not actually
    -- lead to any code generation.  This is supposed to be used alongside
    -- the other monadic expression functions, such as 'eIf'.
    eSubExp ::
      (MonadBuilder m) =>
      SubExp ->
      m (Exp (Rep m))
    eSubExp :: forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m)))
    -> (SubExp -> Exp (Rep m)) -> SubExp -> m (Exp (Rep m))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m))
    -> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> BasicOp
    SubExp

    -- | Treat a parameter as a monadic expression.
    eParam ::
      (MonadBuilder m) =>
      Param t ->
      m (Exp (Rep m))
    eParam :: forall (m :: * -> *) t.
    MonadBuilder m =>
    Param t -> m (Exp (Rep m))
    eParam = SubExp -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp (SubExp -> m (Exp (Rep m)))
    -> (Param t -> SubExp) -> Param t -> m (Exp (Rep m))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> SubExp
    Var (VName -> SubExp) -> (Param t -> VName) -> Param t -> SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param t -> VName
    forall dec. Param dec -> VName
    paramName

    removeRedundantScrutinees :: [SubExp] -> [Case b] -> ([SubExp], [Case b])
    removeRedundantScrutinees :: forall b. [SubExp] -> [Case b] -> ([SubExp], [Case b])
    removeRedundantScrutinees [SubExp]
    ses [Case b]
    cases =
      let ([SubExp]
    ses', [[Maybe PrimValue]]
    vs) =
            [(SubExp, [Maybe PrimValue])] -> ([SubExp], [[Maybe PrimValue]])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(SubExp, [Maybe PrimValue])] -> ([SubExp], [[Maybe PrimValue]]))
    -> [(SubExp, [Maybe PrimValue])] -> ([SubExp], [[Maybe PrimValue]])
    forall a b. (a -> b) -> a -> b
    $ ((SubExp, [Maybe PrimValue]) -> Bool)
    -> [(SubExp, [Maybe PrimValue])] -> [(SubExp, [Maybe PrimValue])]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (SubExp, [Maybe PrimValue]) -> Bool
    forall {a}. (a, [Maybe PrimValue]) -> Bool
    interesting ([(SubExp, [Maybe PrimValue])] -> [(SubExp, [Maybe PrimValue])])
    -> [(SubExp, [Maybe PrimValue])] -> [(SubExp, [Maybe PrimValue])]
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> [[Maybe PrimValue]] -> [(SubExp, [Maybe PrimValue])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    ses ([[Maybe PrimValue]] -> [(SubExp, [Maybe PrimValue])])
    -> [[Maybe PrimValue]] -> [(SubExp, [Maybe PrimValue])]
    forall a b. (a -> b) -> a -> b
    $ [[Maybe PrimValue]] -> [[Maybe PrimValue]]
    forall a. [[a]] -> [[a]]
    L.transpose ((Case b -> [Maybe PrimValue]) -> [Case b] -> [[Maybe PrimValue]]
    forall a b. (a -> b) -> [a] -> [b]
    map Case b -> [Maybe PrimValue]
    forall body. Case body -> [Maybe PrimValue]
    casePat [Case b]
    cases)
       in ([SubExp]
    ses', ([Maybe PrimValue] -> b -> Case b)
    -> [[Maybe PrimValue]] -> [b] -> [Case b]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith [Maybe PrimValue] -> b -> Case b
    forall body. [Maybe PrimValue] -> body -> Case body
    Case ([[Maybe PrimValue]] -> [[Maybe PrimValue]]
    forall a. [[a]] -> [[a]]
    L.transpose [[Maybe PrimValue]]
    vs) ([b] -> [Case b]) -> [b] -> [Case b]
    forall a b. (a -> b) -> a -> b
    $ (Case b -> b) -> [Case b] -> [b]
    forall a b. (a -> b) -> [a] -> [b]
    map Case b -> b
    forall body. Case body -> body
    caseBody [Case b]
    cases)
      where
        interesting :: (a, [Maybe PrimValue]) -> Bool
    interesting = (Maybe PrimValue -> Bool) -> [Maybe PrimValue] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (Maybe PrimValue -> Maybe PrimValue -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Maybe PrimValue
    forall a. Maybe a
    Nothing) ([Maybe PrimValue] -> Bool)
    -> ((a, [Maybe PrimValue]) -> [Maybe PrimValue])
    -> (a, [Maybe PrimValue])
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (a, [Maybe PrimValue]) -> [Maybe PrimValue]
    forall a b. (a, b) -> b
    snd

    -- | As 'eMatch', but an 'MatchSort' can be given.
    eMatch' ::
      (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
      [SubExp] ->
      [Case (m (Body (Rep m)))] ->
      m (Body (Rep m)) ->
      MatchSort ->
      m (Exp (Rep m))
    eMatch' :: forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    [SubExp]
    -> [Case (m (Body (Rep m)))]
    -> m (Body (Rep m))
    -> MatchSort
    -> m (Exp (Rep m))
    eMatch' [SubExp]
    ses [Case (m (Body (Rep m)))]
    cases_m m (Body (Rep m))
    defbody_m MatchSort
    sort = do
      cases <- (Case (m (Body (Rep m))) -> m (Case (Body (Rep m))))
    -> [Case (m (Body (Rep m)))] -> m [Case (Body (Rep m))]
    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 ((m (Body (Rep m)) -> m (Body (Rep m)))
    -> Case (m (Body (Rep m))) -> m (Case (Body (Rep m)))
    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) -> Case a -> f (Case b)
    traverse m (Body (Rep m)) -> m (Body (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    m (Body (Rep m)) -> m (Body (Rep m))
    insertStmsM) [Case (m (Body (Rep m)))]
    cases_m
      defbody <- insertStmsM defbody_m
      ts <-
        L.foldl' generaliseExtTypes
          <$> bodyExtType defbody
          <*> mapM (bodyExtType . caseBody) cases
      cases' <- mapM (traverse $ addContextForBranch ts) cases
      defbody' <- addContextForBranch ts defbody
      let ts' = Int -> ExtType -> [ExtType]
    forall a. Int -> a -> [a]
    replicate (Set Int -> Int
    forall a. Set a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([ExtType] -> Set Int
    forall u. [TypeBase (ShapeBase (Ext SubExp)) u] -> Set Int
    shapeContext [ExtType]
    ts)) (PrimType -> ExtType
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64) [ExtType] -> [ExtType] -> [ExtType]
    forall a. [a] -> [a] -> [a]
    ++ [ExtType]
    ts
          (ses', cases'') = removeRedundantScrutinees ses cases'
      pure $ Match ses' cases'' defbody' $ MatchDec ts' sort
      where
        addContextForBranch :: [TypeBase (ShapeBase (Ext SubExp)) u]
    -> Body (Rep m) -> m (Body (Rep m))
    addContextForBranch [TypeBase (ShapeBase (Ext SubExp)) u]
    ts (Body BodyDec (Rep m)
    _ Stms (Rep m)
    stms Result
    val_res) = do
          body_ts <- ExtendedScope (Rep m) m [Type] -> Scope (Rep m) -> m [Type]
    forall rep (m :: * -> *) a.
    ExtendedScope rep m a -> Scope rep -> m a
    extendedScope ((SubExpRes -> ExtendedScope (Rep m) m Type)
    -> Result -> ExtendedScope (Rep m) m [Type]
    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 (Rep m) m Type
    forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
    subExpResType Result
    val_res) Scope (Rep m)
    stmsscope
          let ctx_res =
                ((Int, SubExp) -> SubExp) -> [(Int, SubExp)] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int, SubExp) -> SubExp
    forall a b. (a, b) -> b
    snd ([(Int, SubExp)] -> [SubExp]) -> [(Int, SubExp)] -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ ((Int, SubExp) -> Int) -> [(Int, SubExp)] -> [(Int, SubExp)]
    forall b a. Ord b => (a -> b) -> [a] -> [a]
    L.sortOn (Int, SubExp) -> Int
    forall a b. (a, b) -> a
    fst ([(Int, SubExp)] -> [(Int, SubExp)])
    -> [(Int, SubExp)] -> [(Int, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ Map Int SubExp -> [(Int, SubExp)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map Int SubExp -> [(Int, SubExp)])
    -> Map Int SubExp -> [(Int, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ [TypeBase (ShapeBase (Ext SubExp)) u] -> [Type] -> Map Int SubExp
    forall u u1.
    [TypeBase (ShapeBase (Ext SubExp)) u]
    -> [TypeBase Shape u1] -> Map Int SubExp
    shapeExtMapping [TypeBase (ShapeBase (Ext SubExp)) u]
    ts [Type]
    body_ts
          mkBodyM stms $ subExpsRes ctx_res ++ val_res
          where
            stmsscope :: Scope (Rep m)
    stmsscope = Stms (Rep m) -> Scope (Rep m)
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stms (Rep m)
    stms

    -- | Construct a 'Match' expression.  The main convenience here is
    -- that the existential context of the return type is automatically
    -- deduced, and the necessary elements added to the branches.
    eMatch ::
      (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
      [SubExp] ->
      [Case (m (Body (Rep m)))] ->
      m (Body (Rep m)) ->
      m (Exp (Rep m))
    eMatch :: 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 [Case (m (Body (Rep m)))]
    cases_m m (Body (Rep m))
    defbody_m = [SubExp]
    -> [Case (m (Body (Rep m)))]
    -> m (Body (Rep m))
    -> MatchSort
    -> m (Exp (Rep m))
    forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    [SubExp]
    -> [Case (m (Body (Rep m)))]
    -> m (Body (Rep m))
    -> MatchSort
    -> m (Exp (Rep m))
    eMatch' [SubExp]
    ses [Case (m (Body (Rep m)))]
    cases_m m (Body (Rep m))
    defbody_m MatchSort
    MatchNormal

    -- | Construct a 'Match' modelling an if-expression from a monadic
    -- condition and monadic branches.  'eBody' might be convenient for
    -- constructing the branches.
    eIf ::
      (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
      m (Exp (Rep m)) ->
      m (Body (Rep m)) ->
      m (Body (Rep m)) ->
      m (Exp (Rep m))
    eIf :: 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 m (Exp (Rep m))
    ce m (Body (Rep m))
    te m (Body (Rep m))
    fe = m (Exp (Rep m))
    -> m (Body (Rep m))
    -> m (Body (Rep m))
    -> MatchSort
    -> m (Exp (Rep m))
    forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    m (Exp (Rep m))
    -> m (Body (Rep m))
    -> m (Body (Rep m))
    -> MatchSort
    -> m (Exp (Rep m))
    eIf' m (Exp (Rep m))
    ce m (Body (Rep m))
    te m (Body (Rep m))
    fe MatchSort
    MatchNormal

    -- | As 'eIf', but an 'MatchSort' can be given.
    eIf' ::
      (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
      m (Exp (Rep m)) ->
      m (Body (Rep m)) ->
      m (Body (Rep m)) ->
      MatchSort ->
      m (Exp (Rep m))
    eIf' :: forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    m (Exp (Rep m))
    -> m (Body (Rep m))
    -> m (Body (Rep m))
    -> MatchSort
    -> m (Exp (Rep m))
    eIf' m (Exp (Rep m))
    ce m (Body (Rep m))
    te m (Body (Rep m))
    fe MatchSort
    if_sort = do
      ce' <- String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "cond" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    ce
      eMatch' [ce'] [Case [Just $ BoolValue True] te] fe if_sort

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

    -- | Construct a v'BinOp' expression with the given operator.
    eBinOp ::
      (MonadBuilder m) =>
      BinOp ->
      m (Exp (Rep m)) ->
      m (Exp (Rep m)) ->
      m (Exp (Rep m))
    eBinOp :: forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eBinOp BinOp
    op m (Exp (Rep m))
    x m (Exp (Rep m))
    y = do
      x' <- String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    x
      y' <- letSubExp "y" =<< y
      pure $ BasicOp $ BinOp op x' y'

    -- | Construct a v'UnOp' expression with the given operator.
    eUnOp ::
      (MonadBuilder m) =>
      UnOp ->
      m (Exp (Rep m)) ->
      m (Exp (Rep m))
    eUnOp :: forall (m :: * -> *).
    MonadBuilder m =>
    UnOp -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eUnOp UnOp
    op m (Exp (Rep m))
    x = BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m))
    -> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . UnOp -> SubExp -> BasicOp
    UnOp UnOp
    op (SubExp -> Exp (Rep m)) -> m SubExp -> m (Exp (Rep m))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    x)

    -- | Construct a v'CmpOp' expression with the given comparison.
    eCmpOp ::
      (MonadBuilder m) =>
      CmpOp ->
      m (Exp (Rep m)) ->
      m (Exp (Rep m)) ->
      m (Exp (Rep m))
    eCmpOp :: forall (m :: * -> *).
    MonadBuilder m =>
    CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eCmpOp CmpOp
    op m (Exp (Rep m))
    x m (Exp (Rep m))
    y = do
      x' <- String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    x
      y' <- letSubExp "y" =<< y
      pure $ BasicOp $ CmpOp op x' y'

    -- | Construct a v'ConvOp' expression with the given conversion.
    eConvOp ::
      (MonadBuilder m) =>
      ConvOp ->
      m (Exp (Rep m)) ->
      m (Exp (Rep m))
    eConvOp :: forall (m :: * -> *).
    MonadBuilder m =>
    ConvOp -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eConvOp ConvOp
    op m (Exp (Rep m))
    x = do
      x' <- String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    x
      pure $ BasicOp $ ConvOp op x'

    -- | Construct a 'SSignum' expression.  Fails if the provided
    -- expression is not of integer type.
    eSignum ::
      (MonadBuilder m) =>
      m (Exp (Rep m)) ->
      m (Exp (Rep m))
    eSignum :: forall (m :: * -> *).
    MonadBuilder m =>
    m (Exp (Rep m)) -> m (Exp (Rep m))
    eSignum m (Exp (Rep m))
    em = do
      e <- m (Exp (Rep m))
    em
      e' <- letSubExp "signum_arg" e
      t <- subExpType e'
      case t of
        Prim (IntType IntType
    int_t) ->
          Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ UnOp -> SubExp -> BasicOp
    UnOp (IntType -> UnOp
    SSignum IntType
    int_t) SubExp
    e'
        Type
    _ ->
          String -> m (Exp (Rep m))
    forall a. HasCallStack => String -> a
    error (String -> m (Exp (Rep m))) -> String -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ String
    "eSignum: operand " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Exp (Rep m) -> String
    forall a. Pretty a => a -> String
    prettyString Exp (Rep m)
    e String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    " has invalid type."

    -- | Copy a value.
    eCopy ::
      (MonadBuilder m) =>
      m (Exp (Rep m)) ->
      m (Exp (Rep m))
    eCopy :: forall (m :: * -> *).
    MonadBuilder m =>
    m (Exp (Rep m)) -> m (Exp (Rep m))
    eCopy m (Exp (Rep m))
    e = BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m))
    -> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Shape -> SubExp -> BasicOp
    Replicate Shape
    forall a. Monoid a => a
    mempty (SubExp -> Exp (Rep m)) -> m SubExp -> m (Exp (Rep m))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "copy_arg" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    e)

    -- | Construct a body from expressions.  If multiple expressions are
    -- provided, their results will be concatenated in order and returned
    -- as the result.
    --
    -- /Beware/: this will not produce correct code if the type of the
    -- body would be existential.  That is, the type of the results being
    -- returned should be invariant to the body.
    eBody ::
      (MonadBuilder m) =>
      [m (Exp (Rep m))] ->
      m (Body (Rep m))
    eBody :: forall (m :: * -> *).
    MonadBuilder m =>
    [m (Exp (Rep m))] -> m (Body (Rep m))
    eBody [m (Exp (Rep m))]
    es = m Result -> m (Body (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ (m Result -> m (Body (Rep m))) -> m Result -> m (Body (Rep m))
    forall a b. (a -> b) -> a -> b
    $ do
      es' <- [m (Exp (Rep m))] -> m [Exp (Rep m)]
    forall (t :: * -> *) (m :: * -> *) a.
    (Traversable t, Monad m) =>
    t (m a) -> m (t a)
    forall (m :: * -> *) a. Monad m => [m a] -> m [a]
    sequence [m (Exp (Rep m))]
    es
      xs <- mapM (letTupExp "x") es'
      pure $ varsRes $ concat xs

    -- | Bind each lambda parameter to the result of an expression, then
    -- bind the body of the lambda.  The expressions must produce only a
    -- single value each.
    eLambda ::
      (MonadBuilder m) =>
      Lambda (Rep m) ->
      [m (Exp (Rep m))] ->
      m [SubExpRes]
    eLambda :: forall (m :: * -> *).
    MonadBuilder m =>
    Lambda (Rep m) -> [m (Exp (Rep m))] -> m Result
    eLambda Lambda (Rep m)
    lam [m (Exp (Rep m))]
    args = do
      (Param (LParamInfo (Rep m)) -> m (Exp (Rep m)) -> m ())
    -> [Param (LParamInfo (Rep m))] -> [m (Exp (Rep m))] -> m ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ Param (LParamInfo (Rep m)) -> m (Exp (Rep m)) -> m ()
    forall {m :: * -> *} {dec}.
    MonadBuilder m =>
    Param dec -> m (Exp (Rep m)) -> m ()
    bindParam (Lambda (Rep m) -> [Param (LParamInfo (Rep m))]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda (Rep m)
    lam) [m (Exp (Rep m))]
    args
      Body (Rep m) -> m Result
    forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
    bodyBind (Body (Rep m) -> m Result) -> Body (Rep m) -> m Result
    forall a b. (a -> b) -> a -> b
    $ Lambda (Rep m) -> Body (Rep m)
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda (Rep m)
    lam
      where
        bindParam :: Param dec -> m (Exp (Rep m)) -> m ()
    bindParam Param dec
    param m (Exp (Rep m))
    arg = [VName] -> Exp (Rep m) -> m ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [Param dec -> VName
    forall dec. Param dec -> VName
    paramName Param dec
    param] (Exp (Rep m) -> m ()) -> m (Exp (Rep m)) -> m ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< m (Exp (Rep m))
    arg

    -- | @eInBoundsForDim w i@ produces @0 <= i < w@.
    eDimInBounds :: (MonadBuilder m) => m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eDimInBounds :: forall (m :: * -> *).
    MonadBuilder m =>
    m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eDimInBounds m (Exp (Rep m))
    w m (Exp (Rep m))
    i =
      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
        BinOp
    LogAnd
        (CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eCmpOp (IntType -> CmpOp
    CmpSle IntType
    Int64) (SubExp -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0)) m (Exp (Rep m))
    i)
        (CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
    eCmpOp (IntType -> CmpOp
    CmpSlt IntType
    Int64) m (Exp (Rep m))
    i m (Exp (Rep m))
    w)

    -- | Are these indexes out-of-bounds for the array?
    eOutOfBounds ::
      (MonadBuilder m) =>
      VName ->
      [m (Exp (Rep m))] ->
      m (Exp (Rep m))
    eOutOfBounds :: forall (m :: * -> *).
    MonadBuilder m =>
    VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
    eOutOfBounds VName
    arr [m (Exp (Rep m))]
    is = do
      arr_t <- VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
      let ws = Type -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims Type
    arr_t
      is' <- mapM (letSubExp "write_i") =<< sequence is
      let checkDim SubExp
    w SubExp
    i = do
            less_than_zero <-
              String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "less_than_zero" (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
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $
                  CmpOp -> SubExp -> SubExp -> BasicOp
    CmpOp (IntType -> CmpOp
    CmpSlt IntType
    Int64) SubExp
    i (Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64
    0 :: Int64))
            greater_than_size <-
              letSubExp "greater_than_size" $
                BasicOp $
                  CmpOp (CmpSle Int64) w i
            letSubExp "outside_bounds_dim" $
              BasicOp $
                BinOp LogOr less_than_zero greater_than_size
      foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is'

    -- | The array element at this index.  Returns array unmodified if
    -- indexes are null (does not even need to be an array in that case).
    eIndex :: (MonadBuilder m) => VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
    eIndex :: forall (m :: * -> *).
    MonadBuilder m =>
    VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
    eIndex VName
    arr [] = SubExp -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp (SubExp -> m (Exp (Rep m))) -> SubExp -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    arr
    eIndex VName
    arr [m (Exp (Rep m))]
    is = do
      is' <- (m (Exp (Rep m)) -> m SubExp) -> [m (Exp (Rep m))] -> m [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 (String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "i" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<<) [m (Exp (Rep m))]
    is
      arr_t <- lookupType arr
      pure $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix is'

    -- | The last element of the given array.
    eLast :: (MonadBuilder m) => VName -> m (Exp (Rep m))
    eLast :: forall (m :: * -> *). MonadBuilder m => VName -> m (Exp (Rep m))
    eLast VName
    arr = do
      n <- Int -> Type -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    arraySize Int
    0 (Type -> SubExp) -> m Type -> m SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
      nm1 <-
        letSubExp "nm1" . BasicOp $
          BinOp (Sub Int64 OverflowUndef) n (intConst Int64 1)
      eIndex arr [eSubExp nm1]

    -- | Construct an unspecified value of the given type.
    eBlank :: (MonadBuilder m) => Type -> m (Exp (Rep m))
    eBlank :: forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
    eBlank (Prim PrimType
    t) = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ PrimValue -> SubExp
    Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
    forall a b. (a -> b) -> a -> b
    $ PrimType -> PrimValue
    blankPrimValue PrimType
    t
    eBlank (Array PrimType
    t Shape
    shape NoUniqueness
    _) = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ PrimType -> [SubExp] -> BasicOp
    Scratch PrimType
    t ([SubExp] -> BasicOp) -> [SubExp] -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ Shape -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims Shape
    shape
    eBlank Acc {} = String -> m (Exp (Rep m))
    forall a. HasCallStack => String -> a
    error String
    "eBlank: cannot create blank accumulator"
    eBlank Mem {} = String -> m (Exp (Rep m))
    forall a. HasCallStack => String -> a
    error String
    "eBlank: cannot create blank memory"

    -- | Sign-extend to the given integer type.
    asIntS :: (MonadBuilder m) => IntType -> SubExp -> m SubExp
    asIntS :: forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntS = (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
    asInt IntType -> IntType -> ConvOp
    SExt

    -- | Zero-extend to the given integer type.
    asIntZ :: (MonadBuilder m) => IntType -> SubExp -> m SubExp
    asIntZ :: forall (m :: * -> *).
    MonadBuilder m =>
    IntType -> SubExp -> m SubExp
    asIntZ = (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
    asInt IntType -> IntType -> ConvOp
    ZExt

    asInt ::
      (MonadBuilder m) =>
      (IntType -> IntType -> ConvOp) ->
      IntType ->
      SubExp ->
      m SubExp
    asInt :: forall (m :: * -> *).
    MonadBuilder m =>
    (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
    asInt IntType -> IntType -> ConvOp
    ext IntType
    to_it SubExp
    e = do
      e_t <- SubExp -> m Type
    forall t (m :: * -> *). HasScope t m => SubExp -> m Type
    subExpType SubExp
    e
      case e_t of
        Prim (IntType IntType
    from_it)
          | IntType
    to_it IntType -> IntType -> Bool
    forall a. Eq a => a -> a -> Bool
    == IntType
    from_it -> SubExp -> m SubExp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    e
          | Bool
    otherwise -> String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    s (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
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ ConvOp -> SubExp -> BasicOp
    ConvOp (IntType -> IntType -> ConvOp
    ext IntType
    from_it IntType
    to_it) SubExp
    e
        Type
    _ -> String -> m SubExp
    forall a. HasCallStack => String -> a
    error String
    "asInt: wrong type"
      where
        s :: String
    s = case SubExp
    e of
          Var VName
    v -> VName -> String
    baseString VName
    v
          SubExp
    _ -> String
    "to_" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ IntType -> String
    forall a. Pretty a => a -> String
    prettyString IntType
    to_it

    -- | Apply a binary operator to several subexpressions.  A left-fold.
    foldBinOp ::
      (MonadBuilder m) =>
      BinOp ->
      SubExp ->
      [SubExp] ->
      m (Exp (Rep m))
    foldBinOp :: forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    foldBinOp BinOp
    _ SubExp
    ne [] =
      Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp SubExp
    ne
    foldBinOp BinOp
    bop SubExp
    ne (SubExp
    e : [SubExp]
    es) =
      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 BinOp
    bop (Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp SubExp
    e) (BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    foldBinOp BinOp
    bop SubExp
    ne [SubExp]
    es)

    -- | True if all operands are true.
    eAll :: (MonadBuilder m) => [SubExp] -> m (Exp (Rep m))
    eAll :: forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAll [] = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    True
    eAll [SubExp
    x] = SubExp -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    x
    eAll (SubExp
    x : [SubExp]
    xs) = BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    foldBinOp BinOp
    LogAnd SubExp
    x [SubExp]
    xs

    -- | True if any operand is true.
    eAny :: (MonadBuilder m) => [SubExp] -> m (Exp (Rep m))
    eAny :: forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
    eAny [] = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ Bool -> SubExp
    forall v. IsValue v => v -> SubExp
    constant Bool
    False
    eAny [SubExp
    x] = SubExp -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    x
    eAny (SubExp
    x : [SubExp]
    xs) = BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    foldBinOp BinOp
    LogOr SubExp
    x [SubExp]
    xs

    -- | Create a two-parameter lambda whose body applies the given binary
    -- operation to its arguments.  It is assumed that both argument and
    -- result types are the same.  (This assumption should be fixed at
    -- some point.)
    binOpLambda ::
      (MonadBuilder m, Buildable (Rep m)) =>
      BinOp ->
      PrimType ->
      m (Lambda (Rep m))
    binOpLambda :: forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    BinOp -> PrimType -> m (Lambda (Rep m))
    binOpLambda BinOp
    bop PrimType
    t = (SubExp -> SubExp -> BasicOp)
    -> PrimType -> PrimType -> m (Lambda (Rep m))
    forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    (SubExp -> SubExp -> BasicOp)
    -> PrimType -> PrimType -> m (Lambda (Rep m))
    binLambda (BinOp -> SubExp -> SubExp -> BasicOp
    BinOp BinOp
    bop) PrimType
    t PrimType
    t

    -- | As 'binOpLambda', but for t'CmpOp's.
    cmpOpLambda ::
      (MonadBuilder m, Buildable (Rep m)) =>
      CmpOp ->
      m (Lambda (Rep m))
    cmpOpLambda :: forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    CmpOp -> m (Lambda (Rep m))
    cmpOpLambda CmpOp
    cop = (SubExp -> SubExp -> BasicOp)
    -> PrimType -> PrimType -> m (Lambda (Rep m))
    forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    (SubExp -> SubExp -> BasicOp)
    -> PrimType -> PrimType -> m (Lambda (Rep m))
    binLambda (CmpOp -> SubExp -> SubExp -> BasicOp
    CmpOp CmpOp
    cop) (CmpOp -> PrimType
    cmpOpType CmpOp
    cop) PrimType
    Bool

    binLambda ::
      (MonadBuilder m, Buildable (Rep m)) =>
      (SubExp -> SubExp -> BasicOp) ->
      PrimType ->
      PrimType ->
      m (Lambda (Rep m))
    binLambda :: forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    (SubExp -> SubExp -> BasicOp)
    -> PrimType -> PrimType -> m (Lambda (Rep m))
    binLambda SubExp -> SubExp -> BasicOp
    bop PrimType
    arg_t PrimType
    ret_t = do
      x <- String -> m VName
    forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newVName String
    "x"
      y <- newVName "y"
      body <-
        buildBody_ . fmap (pure . subExpRes) $
          letSubExp "binlam_res" $
            BasicOp $
              bop (Var x) (Var y)
      pure
        Lambda
          { lambdaParams =
              [ Param mempty x (Prim arg_t),
                Param mempty y (Prim arg_t)
              ],
            lambdaReturnType = [Prim ret_t],
            lambdaBody = body
          }

    -- | Easily construct a t'Lambda' within a 'MonadBuilder'.  See also
    -- 'runLambdaBuilder'.
    mkLambda ::
      (MonadBuilder m) =>
      [LParam (Rep m)] ->
      m Result ->
      m (Lambda (Rep m))
    mkLambda :: forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [LParam (Rep m)]
    params m Result
    m = do
      (body, ret) <- m (Result, [Type]) -> m (Body (Rep m), [Type])
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m (Result, a) -> m (Body (Rep m), a)
    buildBody (m (Result, [Type]) -> m (Body (Rep m), [Type]))
    -> (m (Result, [Type]) -> m (Result, [Type]))
    -> m (Result, [Type])
    -> m (Body (Rep m), [Type])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Scope (Rep m) -> m (Result, [Type]) -> m (Result, [Type])
    forall a. Scope (Rep m) -> m a -> m a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([LParam (Rep m)] -> Scope (Rep m)
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams [LParam (Rep m)]
    params) (m (Result, [Type]) -> m (Body (Rep m), [Type]))
    -> m (Result, [Type]) -> m (Body (Rep m), [Type])
    forall a b. (a -> b) -> a -> b
    $ do
        res <- m Result
    m
        ret <- mapM subExpResType res
        pure (res, ret)
      pure $ Lambda params ret body

    -- | Slice a full dimension of the given size.
    sliceDim :: SubExp -> DimIndex SubExp
    sliceDim :: SubExp -> DimIndex SubExp
    sliceDim SubExp
    d = SubExp -> SubExp -> SubExp -> DimIndex SubExp
    forall d. d -> d -> d -> DimIndex d
    DimSlice (Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64
    0 :: Int64)) SubExp
    d (Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64
    1 :: Int64))

    -- | @fullSlice t slice@ returns @slice@, but with 'DimSlice's of
    -- entire dimensions appended to the full dimensionality of @t@.  This
    -- function is used to turn incomplete indexing complete, as required
    -- by 'Index'.
    fullSlice :: Type -> [DimIndex SubExp] -> Slice SubExp
    fullSlice :: Type -> [DimIndex SubExp] -> Slice SubExp
    fullSlice Type
    t [DimIndex SubExp]
    slice =
      [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice ([DimIndex SubExp] -> Slice SubExp)
    -> [DimIndex SubExp] -> Slice SubExp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp]
    slice [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
    forall a. [a] -> [a] -> [a]
    ++ (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> DimIndex SubExp
    sliceDim (Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    drop ([DimIndex SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [DimIndex SubExp]
    slice) ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Type -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims Type
    t)

    -- | @ sliceAt t n slice@ returns @slice@ but with 'DimSlice's of the
    -- outer @n@ dimensions prepended, and as many appended as to make it
    -- a full slice.  This is a generalisation of 'fullSlice'.
    sliceAt :: Type -> Int -> [DimIndex SubExp] -> Slice SubExp
    sliceAt :: Type -> Int -> [DimIndex SubExp] -> Slice SubExp
    sliceAt Type
    t Int
    n [DimIndex SubExp]
    slice =
      Type -> [DimIndex SubExp] -> Slice SubExp
    fullSlice Type
    t ([DimIndex SubExp] -> Slice SubExp)
    -> [DimIndex SubExp] -> Slice SubExp
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> DimIndex SubExp
    sliceDim (Int -> [SubExp] -> [SubExp]
    forall a. Int -> [a] -> [a]
    take Int
    n ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Type -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims Type
    t) [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [DimIndex SubExp]
    slice

    -- | Like 'fullSlice', but the dimensions are simply numeric.
    fullSliceNum :: (Num d) => [d] -> [DimIndex d] -> Slice d
    fullSliceNum :: forall d. Num d => [d] -> [DimIndex d] -> Slice d
    fullSliceNum [d]
    dims [DimIndex d]
    slice =
      [DimIndex d] -> Slice d
    forall d. [DimIndex d] -> Slice d
    Slice ([DimIndex d] -> Slice d) -> [DimIndex d] -> Slice d
    forall a b. (a -> b) -> a -> b
    $ [DimIndex d]
    slice [DimIndex d] -> [DimIndex d] -> [DimIndex d]
    forall a. [a] -> [a] -> [a]
    ++ (d -> DimIndex d) -> [d] -> [DimIndex d]
    forall a b. (a -> b) -> [a] -> [b]
    map (\d
    d -> d -> d -> d -> DimIndex d
    forall d. d -> d -> d -> DimIndex d
    DimSlice d
    0 d
    d d
    1) (Int -> [d] -> [d]
    forall a. Int -> [a] -> [a]
    drop ([DimIndex d] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [DimIndex d]
    slice) [d]
    dims)

    -- | Does the slice describe the full size of the array?  The most
    -- obvious such slice is one that 'DimSlice's the full span of every
    -- dimension, but also one that fixes all unit dimensions.
    isFullSlice :: Shape -> Slice SubExp -> Bool
    isFullSlice :: Shape -> Slice SubExp -> Bool
    isFullSlice Shape
    shape Slice SubExp
    slice = [Bool] -> Bool
    forall (t :: * -> *). Foldable t => t Bool -> Bool
    and ([Bool] -> Bool) -> [Bool] -> Bool
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> DimIndex SubExp -> Bool)
    -> [SubExp] -> [DimIndex SubExp] -> [Bool]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith SubExp -> DimIndex SubExp -> Bool
    allOfIt (Shape -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims Shape
    shape) (Slice SubExp -> [DimIndex SubExp]
    forall d. Slice d -> [DimIndex d]
    unSlice Slice SubExp
    slice)
      where
        allOfIt :: SubExp -> DimIndex SubExp -> Bool
    allOfIt (Constant PrimValue
    v) DimFix {} = PrimValue -> Bool
    oneIsh PrimValue
    v
        allOfIt SubExp
    d (DimSlice SubExp
    _ SubExp
    n SubExp
    _) = SubExp
    d SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == SubExp
    n
        allOfIt SubExp
    _ DimIndex SubExp
    _ = Bool
    False

    -- | Conveniently construct a body that contains no bindings.
    resultBody :: (Buildable rep) => [SubExp] -> Body rep
    resultBody :: forall rep. Buildable rep => [SubExp] -> Body rep
    resultBody = Stms rep -> Result -> Body rep
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms rep
    forall a. Monoid a => a
    mempty (Result -> Body rep)
    -> ([SubExp] -> Result) -> [SubExp] -> Body rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [SubExp] -> Result
    subExpsRes

    -- | Conveniently construct a body that contains no bindings - but
    -- this time, monadically!
    resultBodyM :: (MonadBuilder m) => [SubExp] -> m (Body (Rep m))
    resultBodyM :: forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM = Stms (Rep m) -> Result -> m (Body (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    Stms (Rep m) -> Result -> m (Body (Rep m))
    mkBodyM Stms (Rep m)
    forall a. Monoid a => a
    mempty (Result -> m (Body (Rep m)))
    -> ([SubExp] -> Result) -> [SubExp] -> m (Body (Rep m))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [SubExp] -> Result
    subExpsRes

    -- | Evaluate the action, producing a body, then wrap it in all the
    -- bindings it created using 'addStm'.
    insertStmsM ::
      (MonadBuilder m) =>
      m (Body (Rep m)) ->
      m (Body (Rep m))
    insertStmsM :: forall (m :: * -> *).
    MonadBuilder m =>
    m (Body (Rep m)) -> m (Body (Rep m))
    insertStmsM m (Body (Rep m))
    m = do
      (Body _ stms res, otherstms) <- m (Body (Rep m)) -> m (Body (Rep m), Stms (Rep m))
    forall a. m a -> m (a, Stms (Rep m))
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m a -> m (a, Stms (Rep m))
    collectStms m (Body (Rep m))
    m
      mkBodyM (otherstms <> stms) res

    -- | Evaluate an action that produces a 'Result' and an auxiliary
    -- value, then return the body constructed from the 'Result' and any
    -- statements added during the action, along the auxiliary value.
    buildBody ::
      (MonadBuilder m) =>
      m (Result, a) ->
      m (Body (Rep m), a)
    buildBody :: forall (m :: * -> *) a.
    MonadBuilder m =>
    m (Result, a) -> m (Body (Rep m), a)
    buildBody m (Result, a)
    m = do
      ((res, v), stms) <- m (Result, a) -> m ((Result, a), Stms (Rep m))
    forall a. m a -> m (a, Stms (Rep m))
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m a -> m (a, Stms (Rep m))
    collectStms m (Result, a)
    m
      body <- mkBodyM stms res
      pure (body, v)

    -- | As 'buildBody', but there is no auxiliary value.
    buildBody_ ::
      (MonadBuilder m) =>
      m Result ->
      m (Body (Rep m))
    buildBody_ :: forall (m :: * -> *).
    MonadBuilder m =>
    m Result -> m (Body (Rep m))
    buildBody_ m Result
    m = (Body (Rep m), ()) -> Body (Rep m)
    forall a b. (a, b) -> a
    fst ((Body (Rep m), ()) -> Body (Rep m))
    -> m (Body (Rep m), ()) -> m (Body (Rep m))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> m (Result, ()) -> m (Body (Rep m), ())
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m (Result, a) -> m (Body (Rep m), a)
    buildBody ((,()) (Result -> (Result, ())) -> m Result -> m (Result, ())
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> m Result
    m)

    -- | Change that result where evaluation of the body would stop.  Also
    -- change type annotations at branches.
    mapResult ::
      (Buildable rep) =>
      (Result -> Body rep) ->
      Body rep ->
      Body rep
    mapResult :: forall rep.
    Buildable rep =>
    (Result -> Body rep) -> Body rep -> Body rep
    mapResult Result -> Body rep
    f (Body BodyDec rep
    _ Stms rep
    stms Result
    res) =
      let Body BodyDec rep
    _ Stms rep
    stms2 Result
    newres = Result -> Body rep
    f Result
    res
       in Stms rep -> Result -> Body rep
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody (Stms rep
    stms Stms rep -> Stms rep -> Stms rep
    forall a. Semigroup a => a -> a -> a
    <> Stms rep
    stms2) Result
    newres

    -- | Instantiate all existential parts dimensions of the given
    -- type, using a monadic action to create the necessary t'SubExp's.
    -- You should call this function within some monad that allows you to
    -- collect the actions performed (say, 'State').
    instantiateShapes ::
      (Monad m) =>
      (Int -> m SubExp) ->
      [TypeBase ExtShape u] ->
      m [TypeBase Shape u]
    instantiateShapes :: forall (m :: * -> *) u.
    Monad m =>
    (Int -> m SubExp)
    -> [TypeBase (ShapeBase (Ext SubExp)) u] -> m [TypeBase Shape u]
    instantiateShapes Int -> m SubExp
    f [TypeBase (ShapeBase (Ext SubExp)) u]
    ts = StateT (Map Int SubExp) m [TypeBase Shape u]
    -> Map Int SubExp -> m [TypeBase Shape u]
    forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
    evalStateT ((TypeBase (ShapeBase (Ext SubExp)) u
     -> StateT (Map Int SubExp) m (TypeBase Shape u))
    -> [TypeBase (ShapeBase (Ext SubExp)) u]
    -> StateT (Map Int SubExp) m [TypeBase Shape u]
    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 (ShapeBase (Ext SubExp)) u
    -> StateT (Map Int SubExp) m (TypeBase Shape u)
    instantiate [TypeBase (ShapeBase (Ext SubExp)) u]
    ts) Map Int SubExp
    forall k a. Map k a
    M.empty
      where
        instantiate :: TypeBase (ShapeBase (Ext SubExp)) u
    -> StateT (Map Int SubExp) m (TypeBase Shape u)
    instantiate TypeBase (ShapeBase (Ext SubExp)) u
    t = do
          shape <- (Ext SubExp -> StateT (Map Int SubExp) m SubExp)
    -> [Ext SubExp] -> StateT (Map Int SubExp) m [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 Ext SubExp -> StateT (Map Int SubExp) m SubExp
    instantiate' ([Ext SubExp] -> StateT (Map Int SubExp) m [SubExp])
    -> [Ext SubExp] -> StateT (Map Int SubExp) m [SubExp]
    forall a b. (a -> b) -> a -> b
    $ ShapeBase (Ext SubExp) -> [Ext SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (ShapeBase (Ext SubExp) -> [Ext SubExp])
    -> ShapeBase (Ext SubExp) -> [Ext SubExp]
    forall a b. (a -> b) -> a -> b
    $ TypeBase (ShapeBase (Ext SubExp)) u -> ShapeBase (Ext SubExp)
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    arrayShape TypeBase (ShapeBase (Ext SubExp)) u
    t
          pure $ t `setArrayShape` Shape shape
        instantiate' :: Ext SubExp -> StateT (Map Int SubExp) m SubExp
    instantiate' (Ext Int
    x) = do
          m <- StateT (Map Int SubExp) m (Map Int SubExp)
    forall s (m :: * -> *). MonadState s m => m s
    get
          case M.lookup x m of
            Just SubExp
    se -> SubExp -> StateT (Map Int SubExp) m SubExp
    forall a. a -> StateT (Map Int SubExp) m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    se
            Maybe SubExp
    Nothing -> do
              se <- m SubExp -> StateT (Map Int SubExp) m SubExp
    forall (m :: * -> *) a.
    Monad m =>
    m a -> StateT (Map Int SubExp) m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m SubExp -> StateT (Map Int SubExp) m SubExp)
    -> m SubExp -> StateT (Map Int SubExp) m SubExp
    forall a b. (a -> b) -> a -> b
    $ Int -> m SubExp
    f Int
    x
              put $ M.insert x se m
              pure se
        instantiate' (Free SubExp
    se) = SubExp -> StateT (Map Int SubExp) m SubExp
    forall a. a -> StateT (Map Int SubExp) m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SubExp
    se

    -- | Like 'instantiateShapes', but obtains names from the provided
    -- list.  If an 'Ext' is out of bounds of this list, the function
    -- fails with 'error'.
    instantiateShapes' :: [VName] -> [TypeBase ExtShape u] -> [TypeBase Shape u]
    instantiateShapes' :: forall u.
    [VName]
    -> [TypeBase (ShapeBase (Ext SubExp)) u] -> [TypeBase Shape u]
    instantiateShapes' [VName]
    names [TypeBase (ShapeBase (Ext SubExp)) u]
    ts =
      -- Carefully ensure that the order of idents we produce corresponds
      -- to their existential index.
      Identity [TypeBase Shape u] -> [TypeBase Shape u]
    forall a. Identity a -> a
    runIdentity (Identity [TypeBase Shape u] -> [TypeBase Shape u])
    -> Identity [TypeBase Shape u] -> [TypeBase Shape u]
    forall a b. (a -> b) -> a -> b
    $ (Int -> Identity SubExp)
    -> [TypeBase (ShapeBase (Ext SubExp)) u]
    -> Identity [TypeBase Shape u]
    forall (m :: * -> *) u.
    Monad m =>
    (Int -> m SubExp)
    -> [TypeBase (ShapeBase (Ext SubExp)) u] -> m [TypeBase Shape u]
    instantiateShapes Int -> Identity SubExp
    instantiate [TypeBase (ShapeBase (Ext SubExp)) u]
    ts
      where
        instantiate :: Int -> Identity SubExp
    instantiate Int
    x =
          case Int -> [VName] -> Maybe VName
    forall int a. Integral int => int -> [a] -> Maybe a
    maybeNth Int
    x [VName]
    names of
            Maybe VName
    Nothing -> String -> Identity SubExp
    forall a. HasCallStack => String -> a
    error (String -> Identity SubExp) -> String -> Identity SubExp
    forall a b. (a -> b) -> a -> b
    $ String
    "instantiateShapes': " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ [VName] -> String
    forall a. Pretty a => a -> String
    prettyString [VName]
    names String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    ", " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show Int
    x
            Just VName
    name -> SubExp -> Identity SubExp
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp -> Identity SubExp) -> SubExp -> Identity SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    name

    -- | Remove existentials by imposing sizes from another type where
    -- needed.
    removeExistentials :: ExtType -> Type -> Type
    removeExistentials :: ExtType -> Type -> Type
    removeExistentials ExtType
    t1 Type
    t2 =
      ExtType
    t1
        ExtType -> [SubExp] -> Type
    forall oldshape u.
    TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
    `setArrayDims` (Ext SubExp -> SubExp -> SubExp)
    -> [Ext SubExp] -> [SubExp] -> [SubExp]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith
          Ext SubExp -> SubExp -> SubExp
    forall {p}. Ext p -> p -> p
    nonExistential
          (ShapeBase (Ext SubExp) -> [Ext SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (ShapeBase (Ext SubExp) -> [Ext SubExp])
    -> ShapeBase (Ext SubExp) -> [Ext SubExp]
    forall a b. (a -> b) -> a -> b
    $ ExtType -> ShapeBase (Ext SubExp)
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    arrayShape ExtType
    t1)
          (Type -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims Type
    t2)
      where
        nonExistential :: Ext p -> p -> p
    nonExistential (Ext Int
    _) p
    dim = p
    dim
        nonExistential (Free p
    dim) p
    _ = p
    dim

    -- | Can be used as the definition of 'mkLetNames' for a 'Buildable'
    -- instance for simple representations.
    simpleMkLetNames ::
      ( ExpDec rep ~ (),
        LetDec rep ~ Type,
        MonadFreshNames m,
        TypedOp (OpC rep),
        HasScope rep m
      ) =>
      [VName] ->
      Exp rep ->
      m (Stm rep)
    simpleMkLetNames :: forall rep (m :: * -> *).
    (ExpDec rep ~ (), LetDec rep ~ Type, MonadFreshNames m,
     TypedOp (OpC rep), HasScope rep m) =>
    [VName] -> Exp rep -> m (Stm rep)
    simpleMkLetNames [VName]
    names Exp rep
    e = do
      et <- Exp rep -> m [ExtType]
    forall rep (m :: * -> *).
    (HasScope rep m, TypedOp (OpC rep)) =>
    Exp rep -> m [ExtType]
    expExtType Exp rep
    e
      let ts = [VName] -> [ExtType] -> [Type]
    forall u.
    [VName]
    -> [TypeBase (ShapeBase (Ext SubExp)) u] -> [TypeBase Shape u]
    instantiateShapes' [VName]
    names [ExtType]
    et
      pure $ Let (Pat $ zipWith PatElem names ts) (defAux ()) e

    -- | Instances of this class can be converted to Futhark expressions
    -- within a 'MonadBuilder'.
    class ToExp a where
      toExp :: (MonadBuilder m) => a -> m (Exp (Rep m))

    instance ToExp SubExp where
      toExp :: forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    toExp = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m)))
    -> (SubExp -> Exp (Rep m)) -> SubExp -> m (Exp (Rep m))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m))
    -> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> BasicOp
    SubExp

    instance ToExp VName where
      toExp :: forall (m :: * -> *). MonadBuilder m => VName -> m (Exp (Rep m))
    toExp = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m)))
    -> (VName -> Exp (Rep m)) -> VName -> m (Exp (Rep m))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m))
    -> (VName -> BasicOp) -> VName -> Exp (Rep m)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> SubExp
    Var

    -- | A convenient composition of 'letSubExp' and 'toExp'.
    toSubExp :: (MonadBuilder m, ToExp a) => String -> a -> m SubExp
    toSubExp :: forall (m :: * -> *) a.
    (MonadBuilder m, ToExp a) =>
    String -> a -> m SubExp
    toSubExp String
    s a
    e = String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    s (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< a -> m (Exp (Rep m))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *). MonadBuilder m => a -> m (Exp (Rep m))
    toExp a
    e