    -- | This monomorphization module converts a well-typed, polymorphic,
    -- module-free Futhark program into an equivalent monomorphic program.
    --
    -- This pass also does a few other simplifications to make the job of
    -- subsequent passes easier.  Specifically, it does the following:
    --
    -- * Turn operator sections into explicit lambdas.
    --
    -- * Converts applications of intrinsic SOACs into SOAC AST nodes
    --   (Map, Reduce, etc).
    --
    -- * Elide functions that are not reachable from an entry point (this
    --   is a side effect of the monomorphisation algorithm, which uses
    --   the entry points as roots).
    --
    -- * Rewrite BinOp nodes to Apply nodes.
    --
    -- * Replace all size expressions by constants or variables,
    --   complex expressions replaced by variables are calculated in
    --   let binding or replaced by size parameters if in argument.
    --
    -- Note that these changes are unfortunately not visible in the AST
    -- representation.
    module Futhark.Internalise.Monomorphise (transformProg) where

    import Control.Monad
    import Control.Monad.Identity
    import Control.Monad.RWS (MonadReader (..), MonadWriter (..), RWST, asks, runRWST)
    import Control.Monad.State
    import Control.Monad.Writer (Writer, runWriter, runWriterT)
    import Data.Bifunctor
    import Data.Bitraversable
    import Data.Foldable
    import Data.List (partition)
    import Data.List.NonEmpty qualified as NE
    import Data.Map.Strict qualified as M
    import Data.Maybe (isJust, isNothing)
    import Data.Sequence qualified as Seq
    import Data.Set qualified as S
    import Futhark.MonadFreshNames
    import Futhark.Util (nubOrd, topologicalSort)
    import Futhark.Util.Pretty
    import Language.Futhark
    import Language.Futhark.Traversals
    import Language.Futhark.TypeChecker.Types

    i64 :: TypeBase dim als
    i64 :: forall dim als. TypeBase dim als
    i64 = ScalarTypeBase dim als -> TypeBase dim als
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase dim als -> TypeBase dim als)
    -> ScalarTypeBase dim als -> TypeBase dim als
    forall a b. (a -> b) -> a -> b
    $ PrimType -> ScalarTypeBase dim als
    forall dim u. PrimType -> ScalarTypeBase dim u
    Prim (PrimType -> ScalarTypeBase dim als)
    -> PrimType -> ScalarTypeBase dim als
    forall a b. (a -> b) -> a -> b
    $ IntType -> PrimType
    Signed IntType
    Int64

    -- The monomorphization monad reads 'PolyBinding's and writes
    -- 'ValBind's.  The 'TypeParam's in the 'ValBind's can only be size
    -- parameters.
    newtype PolyBinding
      = PolyBinding
          ( Maybe EntryPoint,
            VName,
            [TypeParam],
            [Pat ParamType],
            ResRetType,
            Exp,
            [AttrInfo VName],
            SrcLoc
          )

    -- | To deduplicate size expressions, we want a looser notion of
    -- equality than the strict syntactical equality provided by the Eq
    -- instance on Exp. This newtype wrapper provides such a looser notion
    -- of equality.
    newtype ReplacedExp = ReplacedExp {ReplacedExp -> Exp
    unReplaced :: Exp}
      deriving (Int -> ReplacedExp -> ShowS
    [ReplacedExp] -> ShowS
    ReplacedExp -> [Char]
    (Int -> ReplacedExp -> ShowS)
    -> (ReplacedExp -> [Char])
    -> ([ReplacedExp] -> ShowS)
    -> Show ReplacedExp
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> ReplacedExp -> ShowS
    showsPrec :: Int -> ReplacedExp -> ShowS
    $cshow :: ReplacedExp -> [Char]
    show :: ReplacedExp -> [Char]
    $cshowList :: [ReplacedExp] -> ShowS
    showList :: [ReplacedExp] -> ShowS
    Show)

    instance Pretty ReplacedExp where
      pretty :: forall ann. ReplacedExp -> Doc ann
    pretty (ReplacedExp Exp
    e) = Exp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Exp -> Doc ann
    pretty Exp
    e

    instance Eq ReplacedExp where
      ReplacedExp Exp
    e1 == :: ReplacedExp -> ReplacedExp -> Bool
    == ReplacedExp Exp
    e2
        | Just [(Exp, Exp)]
    es <- Exp -> Exp -> Maybe [(Exp, Exp)]
    similarExps Exp
    e1 Exp
    e2 =
            ((Exp, Exp) -> Bool) -> [(Exp, Exp)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((ReplacedExp -> ReplacedExp -> Bool)
    -> (ReplacedExp, ReplacedExp) -> Bool
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ReplacedExp -> ReplacedExp -> Bool
    forall a. Eq a => a -> a -> Bool
    (==) ((ReplacedExp, ReplacedExp) -> Bool)
    -> ((Exp, Exp) -> (ReplacedExp, ReplacedExp)) -> (Exp, Exp) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Exp -> ReplacedExp)
    -> (Exp -> ReplacedExp) -> (Exp, Exp) -> (ReplacedExp, ReplacedExp)
    forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
    forall (p :: * -> * -> *) a b c d.
    Bifunctor p =>
    (a -> b) -> (c -> d) -> p a c -> p b d
    bimap Exp -> ReplacedExp
    ReplacedExp Exp -> ReplacedExp
    ReplacedExp) [(Exp, Exp)]
    es
      ReplacedExp
    _ == ReplacedExp
    _ = Bool
    False

    type ExpReplacements = [(ReplacedExp, VName)]

    canCalculate :: S.Set VName -> ExpReplacements -> ExpReplacements
    canCalculate :: Set VName -> ExpReplacements -> ExpReplacements
    canCalculate Set VName
    scope ExpReplacements
    mapping = do
      ((ReplacedExp, VName) -> Bool)
    -> ExpReplacements -> ExpReplacements
    forall a. (a -> Bool) -> [a] -> [a]
    filter
        ( (Set VName -> Set VName -> Bool
    forall a. Ord a => Set a -> Set a -> Bool
    `S.isSubsetOf` Set VName
    scope)
            (Set VName -> Bool)
    -> ((ReplacedExp, VName) -> Set VName)
    -> (ReplacedExp, VName)
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> Bool) -> Set VName -> Set VName
    forall a. (a -> Bool) -> Set a -> Set a
    S.filter VName -> Bool
    notIntrisic
            (Set VName -> Set VName)
    -> ((ReplacedExp, VName) -> Set VName)
    -> (ReplacedExp, VName)
    -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . FV -> Set VName
    fvVars
            (FV -> Set VName)
    -> ((ReplacedExp, VName) -> FV)
    -> (ReplacedExp, VName)
    -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Exp -> FV
    freeInExp
            (Exp -> FV)
    -> ((ReplacedExp, VName) -> Exp) -> (ReplacedExp, VName) -> FV
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ReplacedExp -> Exp
    unReplaced
            (ReplacedExp -> Exp)
    -> ((ReplacedExp, VName) -> ReplacedExp)
    -> (ReplacedExp, VName)
    -> Exp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (ReplacedExp, VName) -> ReplacedExp
    forall a b. (a, b) -> a
    fst
        )
        ExpReplacements
    mapping
      where
        notIntrisic :: VName -> Bool
    notIntrisic VName
    vn = VName -> Int
    baseTag VName
    vn Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    maxIntrinsicTag

    -- Replace some expressions by a parameter.
    expReplace :: ExpReplacements -> Exp -> Exp
    expReplace :: ExpReplacements -> Exp -> Exp
    expReplace ExpReplacements
    mapping Exp
    e
      | Just VName
    vn <- ReplacedExp -> ExpReplacements -> Maybe VName
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup (Exp -> ReplacedExp
    ReplacedExp Exp
    e) ExpReplacements
    mapping =
          QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    vn) (StructType -> Info StructType
    forall a. a -> Info a
    Info (StructType -> Info StructType) -> StructType -> Info StructType
    forall a b. (a -> b) -> a -> b
    $ Exp -> StructType
    typeOf Exp
    e) (Exp -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Exp
    e)
    expReplace ExpReplacements
    mapping Exp
    e = Identity Exp -> Exp
    forall a. Identity a -> a
    runIdentity (Identity Exp -> Exp) -> Identity Exp -> Exp
    forall a b. (a -> b) -> a -> b
    $ ASTMapper Identity -> Exp -> Identity Exp
    forall x (m :: * -> *).
    (ASTMappable x, Monad m) =>
    ASTMapper m -> x -> m x
    forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
    astMap ASTMapper Identity
    mapper Exp
    e
      where
        mapper :: ASTMapper Identity
    mapper = ASTMapper Identity
    forall (m :: * -> *). Monad m => ASTMapper m
    identityMapper {mapOnExp = pure . expReplace mapping}

    -- Construct an Assert expression that checks that the names (values)
    -- in the mapping have the same value as the expression they
    -- represent.  This is injected into entry points, where we cannot
    -- otherwise trust the input.  XXX: the error message generated from
    -- this is not great; we should rework it eventually.
    entryAssert :: ExpReplacements -> Exp -> Exp
    entryAssert :: ExpReplacements -> Exp -> Exp
    entryAssert [] Exp
    body = Exp
    body
    entryAssert ((ReplacedExp, VName)
    x : ExpReplacements
    xs) Exp
    body =
      Exp -> Exp -> Info Text -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
    Assert ((Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Exp -> Exp -> Exp
    logAnd ((ReplacedExp, VName) -> Exp
    cmpExp (ReplacedExp, VName)
    x) ([Exp] -> Exp) -> [Exp] -> Exp
    forall a b. (a -> b) -> a -> b
    $ ((ReplacedExp, VName) -> Exp) -> ExpReplacements -> [Exp]
    forall a b. (a -> b) -> [a] -> [b]
    map (ReplacedExp, VName) -> Exp
    cmpExp ExpReplacements
    xs) Exp
    body Info Text
    errmsg (Exp -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Exp
    body)
      where
        errmsg :: Info Text
    errmsg = Text -> Info Text
    forall a. a -> Info a
    Info Text
    "entry point arguments have invalid sizes."
        bool :: TypeBase dim u
    bool = ScalarTypeBase dim u -> TypeBase dim u
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase dim u -> TypeBase dim u)
    -> ScalarTypeBase dim u -> TypeBase dim u
    forall a b. (a -> b) -> a -> b
    $ PrimType -> ScalarTypeBase dim u
    forall dim u. PrimType -> ScalarTypeBase dim u
    Prim PrimType
    Bool
        opt :: StructType
    opt = [ParamType] -> RetTypeBase Exp Uniqueness -> StructType
    foldFunType [ParamType
    forall dim als. TypeBase dim als
    bool, ParamType
    forall dim als. TypeBase dim als
    bool] (RetTypeBase Exp Uniqueness -> StructType)
    -> RetTypeBase Exp Uniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ [VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [] TypeBase Exp Uniqueness
    forall dim als. TypeBase dim als
    bool
        andop :: Exp
    andop = QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName (Name -> VName
    intrinsicVar Name
    "&&")) (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    opt) SrcLoc
    forall a. Monoid a => a
    mempty
        eqop :: Exp
    eqop = QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName (Name -> VName
    intrinsicVar Name
    "==")) (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    opt) SrcLoc
    forall a. Monoid a => a
    mempty
        logAnd :: Exp -> Exp -> Exp
    logAnd Exp
    x' Exp
    y =
          Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply Exp
    andop [(Maybe VName
    forall a. Maybe a
    Nothing, Exp
    x'), (Maybe VName
    forall a. Maybe a
    Nothing, Exp
    y)] (AppRes -> Exp) -> AppRes -> Exp
    forall a b. (a -> b) -> a -> b
    $
            StructType -> [VName] -> AppRes
    AppRes StructType
    forall dim als. TypeBase dim als
    bool []
        cmpExp :: (ReplacedExp, VName) -> Exp
    cmpExp (ReplacedExp Exp
    x', VName
    y) =
          Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply Exp
    eqop [(Maybe VName
    forall a. Maybe a
    Nothing, Exp
    x'), (Maybe VName
    forall a. Maybe a
    Nothing, Exp
    y')] (AppRes -> Exp) -> AppRes -> Exp
    forall a b. (a -> b) -> a -> b
    $
            StructType -> [VName] -> AppRes
    AppRes StructType
    forall dim als. TypeBase dim als
    bool []
          where
            y' :: Exp
    y' = QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    y) (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    forall dim als. TypeBase dim als
    i64) SrcLoc
    forall a. Monoid a => a
    mempty

    -- Monomorphization environment mapping names of polymorphic functions
    -- to a representation of their corresponding function bindings.
    data Env = Env
      { Env -> Map VName PolyBinding
    envPolyBindings :: M.Map VName PolyBinding,
        Env -> Set VName
    envScope :: S.Set VName,
        Env -> Set VName
    envGlobalScope :: S.Set VName,
        Env -> ExpReplacements
    envParametrized :: ExpReplacements
      }

    instance Semigroup Env where
      Env Map VName PolyBinding
    pb1 Set VName
    sc1 Set VName
    gs1 ExpReplacements
    pr1 <> :: Env -> Env -> Env
    <> Env Map VName PolyBinding
    pb2 Set VName
    sc2 Set VName
    gs2 ExpReplacements
    pr2 = Map VName PolyBinding
    -> Set VName -> Set VName -> ExpReplacements -> Env
    Env (Map VName PolyBinding
    pb1 Map VName PolyBinding
    -> Map VName PolyBinding -> Map VName PolyBinding
    forall a. Semigroup a => a -> a -> a
    <> Map VName PolyBinding
    pb2) (Set VName
    sc1 Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> Set VName
    sc2) (Set VName
    gs1 Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> Set VName
    gs2) (ExpReplacements
    pr1 ExpReplacements -> ExpReplacements -> ExpReplacements
    forall a. Semigroup a => a -> a -> a
    <> ExpReplacements
    pr2)

    instance Monoid Env where
      mempty :: Env
    mempty = Map VName PolyBinding
    -> Set VName -> Set VName -> ExpReplacements -> Env
    Env Map VName PolyBinding
    forall a. Monoid a => a
    mempty Set VName
    forall a. Monoid a => a
    mempty Set VName
    forall a. Monoid a => a
    mempty ExpReplacements
    forall a. Monoid a => a
    mempty

    localEnv :: Env -> MonoM a -> MonoM a
    localEnv :: forall a. Env -> MonoM a -> MonoM a
    localEnv Env
    env = (Env -> Env) -> MonoM a -> MonoM a
    forall a. (Env -> Env) -> MonoM a -> MonoM a
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
    local (Env
    env <>)

    isolateNormalisation :: MonoM a -> MonoM a
    isolateNormalisation :: forall a. MonoM a -> MonoM a
    isolateNormalisation MonoM a
    m = do
      ExpReplacements
    prevRepl <- MonoM ExpReplacements
    forall s (m :: * -> *). MonadState s m => m s
    get
      ExpReplacements -> MonoM ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put ExpReplacements
    forall a. Monoid a => a
    mempty
      a
    ret <- (Env -> Env) -> MonoM a -> MonoM a
    forall a. (Env -> Env) -> MonoM a -> MonoM a
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
    local (\Env
    env -> Env
    env {envScope = mempty, envParametrized = mempty}) MonoM a
    m
      ExpReplacements -> MonoM ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put ExpReplacements
    prevRepl
      a -> MonoM a
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    ret

    -- | These now have monomorphic types in the given action. This is
    -- used to handle shadowing.
    withMono :: [VName] -> MonoM a -> MonoM a
    withMono :: forall a. [VName] -> MonoM a -> MonoM a
    withMono [] = MonoM a -> MonoM a
    forall a. a -> a
    id
    withMono [VName]
    vs = (Env -> Env) -> MonoM a -> MonoM a
    forall a. (Env -> Env) -> MonoM a -> MonoM a
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
    local ((Env -> Env) -> MonoM a -> MonoM a)
    -> (Env -> Env) -> MonoM a -> MonoM a
    forall a b. (a -> b) -> a -> b
    $ \Env
    env ->
      Env
    env {envPolyBindings = M.filterWithKey keep (envPolyBindings env)}
      where
        keep :: VName -> p -> Bool
    keep VName
    v p
    _ = VName
    v VName -> [VName] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` [VName]
    vs

    withArgs :: S.Set VName -> MonoM a -> MonoM a
    withArgs :: forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    args = Env -> MonoM a -> MonoM a
    forall a. Env -> MonoM a -> MonoM a
    localEnv (Env -> MonoM a -> MonoM a) -> Env -> MonoM a -> MonoM a
    forall a b. (a -> b) -> a -> b
    $ Env
    forall a. Monoid a => a
    mempty {envScope = args}

    withParams :: ExpReplacements -> MonoM a -> MonoM a
    withParams :: forall a. ExpReplacements -> MonoM a -> MonoM a
    withParams ExpReplacements
    params = Env -> MonoM a -> MonoM a
    forall a. Env -> MonoM a -> MonoM a
    localEnv (Env -> MonoM a -> MonoM a) -> Env -> MonoM a -> MonoM a
    forall a b. (a -> b) -> a -> b
    $ Env
    forall a. Monoid a => a
    mempty {envParametrized = params}

    -- The monomorphization monad.
    newtype MonoM a
      = MonoM
          ( RWST
              Env
              (Seq.Seq (VName, ValBind))
              (ExpReplacements, VNameSource)
              (State Lifts)
              a
          )
      deriving
        ( (forall a b. (a -> b) -> MonoM a -> MonoM b)
    -> (forall a b. a -> MonoM b -> MonoM a) -> Functor MonoM
    forall a b. a -> MonoM b -> MonoM a
    forall a b. (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *).
    (forall a b. (a -> b) -> f a -> f b)
    -> (forall a b. a -> f b -> f a) -> Functor f
    $cfmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
    fmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
    $c<$ :: forall a b. a -> MonoM b -> MonoM a
    <$ :: forall a b. a -> MonoM b -> MonoM a
    Functor,
          Functor MonoM
    Functor MonoM =>
    (forall a. a -> MonoM a)
    -> (forall a b. MonoM (a -> b) -> MonoM a -> MonoM b)
    -> (forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c)
    -> (forall a b. MonoM a -> MonoM b -> MonoM b)
    -> (forall a b. MonoM a -> MonoM b -> MonoM a)
    -> Applicative MonoM
    forall a. a -> MonoM a
    forall a b. MonoM a -> MonoM b -> MonoM a
    forall a b. MonoM a -> MonoM b -> MonoM b
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
    forall (f :: * -> *).
    Functor f =>
    (forall a. a -> f a)
    -> (forall a b. f (a -> b) -> f a -> f b)
    -> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
    -> (forall a b. f a -> f b -> f b)
    -> (forall a b. f a -> f b -> f a)
    -> Applicative f
    $cpure :: forall a. a -> MonoM a
    pure :: forall a. a -> MonoM a
    $c<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    <*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    $cliftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
    liftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
    $c*> :: forall a b. MonoM a -> MonoM b -> MonoM b
    *> :: forall a b. MonoM a -> MonoM b -> MonoM b
    $c<* :: forall a b. MonoM a -> MonoM b -> MonoM a
    <* :: forall a b. MonoM a -> MonoM b -> MonoM a
    Applicative,
          Applicative MonoM
    Applicative MonoM =>
    (forall a b. MonoM a -> (a -> MonoM b) -> MonoM b)
    -> (forall a b. MonoM a -> MonoM b -> MonoM b)
    -> (forall a. a -> MonoM a)
    -> Monad MonoM
    forall a. a -> MonoM a
    forall a b. MonoM a -> MonoM b -> MonoM b
    forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
    forall (m :: * -> *).
    Applicative m =>
    (forall a b. m a -> (a -> m b) -> m b)
    -> (forall a b. m a -> m b -> m b)
    -> (forall a. a -> m a)
    -> Monad m
    $c>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
    >>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
    $c>> :: forall a b. MonoM a -> MonoM b -> MonoM b
    >> :: forall a b. MonoM a -> MonoM b -> MonoM b
    $creturn :: forall a. a -> MonoM a
    return :: forall a. a -> MonoM a
    Monad,
          MonadReader Env,
          MonadWriter (Seq.Seq (VName, ValBind))
        )

    instance MonadFreshNames MonoM where
      getNameSource :: MonoM VNameSource
    getNameSource = RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      VNameSource
    -> MonoM VNameSource
    forall a.
    RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> MonoM a
    MonoM (RWST
       Env
       (Seq (VName, ValBind))
       (ExpReplacements, VNameSource)
       (State Lifts)
       VNameSource
     -> MonoM VNameSource)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         VNameSource
    -> MonoM VNameSource
    forall a b. (a -> b) -> a -> b
    $ ((ExpReplacements, VNameSource) -> VNameSource)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         VNameSource
    forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
    gets (ExpReplacements, VNameSource) -> VNameSource
    forall a b. (a, b) -> b
    snd
      putNameSource :: VNameSource -> MonoM ()
    putNameSource = RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ()
    -> MonoM ()
    forall a.
    RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> MonoM a
    MonoM (RWST
       Env
       (Seq (VName, ValBind))
       (ExpReplacements, VNameSource)
       (State Lifts)
       ()
     -> MonoM ())
    -> (VNameSource
        -> RWST
             Env
             (Seq (VName, ValBind))
             (ExpReplacements, VNameSource)
             (State Lifts)
             ())
    -> VNameSource
    -> MonoM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify (((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
     -> RWST
          Env
          (Seq (VName, ValBind))
          (ExpReplacements, VNameSource)
          (State Lifts)
          ())
    -> (VNameSource
        -> (ExpReplacements, VNameSource)
        -> (ExpReplacements, VNameSource))
    -> VNameSource
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VNameSource -> VNameSource)
    -> (ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource)
    forall b c a. (b -> c) -> (a, b) -> (a, c)
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second ((VNameSource -> VNameSource)
     -> (ExpReplacements, VNameSource)
     -> (ExpReplacements, VNameSource))
    -> (VNameSource -> VNameSource -> VNameSource)
    -> VNameSource
    -> (ExpReplacements, VNameSource)
    -> (ExpReplacements, VNameSource)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> VNameSource -> VNameSource
    forall a b. a -> b -> a
    const

    instance MonadState ExpReplacements MonoM where
      get :: MonoM ExpReplacements
    get = RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ExpReplacements
    -> MonoM ExpReplacements
    forall a.
    RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> MonoM a
    MonoM (RWST
       Env
       (Seq (VName, ValBind))
       (ExpReplacements, VNameSource)
       (State Lifts)
       ExpReplacements
     -> MonoM ExpReplacements)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ExpReplacements
    -> MonoM ExpReplacements
    forall a b. (a -> b) -> a -> b
    $ ((ExpReplacements, VNameSource) -> ExpReplacements)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ExpReplacements
    forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
    gets (ExpReplacements, VNameSource) -> ExpReplacements
    forall a b. (a, b) -> a
    fst
      put :: ExpReplacements -> MonoM ()
    put = RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ()
    -> MonoM ()
    forall a.
    RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> MonoM a
    MonoM (RWST
       Env
       (Seq (VName, ValBind))
       (ExpReplacements, VNameSource)
       (State Lifts)
       ()
     -> MonoM ())
    -> (ExpReplacements
        -> RWST
             Env
             (Seq (VName, ValBind))
             (ExpReplacements, VNameSource)
             (State Lifts)
             ())
    -> ExpReplacements
    -> MonoM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify (((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
     -> RWST
          Env
          (Seq (VName, ValBind))
          (ExpReplacements, VNameSource)
          (State Lifts)
          ())
    -> (ExpReplacements
        -> (ExpReplacements, VNameSource)
        -> (ExpReplacements, VNameSource))
    -> ExpReplacements
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (ExpReplacements -> ExpReplacements)
    -> (ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource)
    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 ((ExpReplacements -> ExpReplacements)
     -> (ExpReplacements, VNameSource)
     -> (ExpReplacements, VNameSource))
    -> (ExpReplacements -> ExpReplacements -> ExpReplacements)
    -> ExpReplacements
    -> (ExpReplacements, VNameSource)
    -> (ExpReplacements, VNameSource)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ExpReplacements -> ExpReplacements -> ExpReplacements
    forall a b. a -> b -> a
    const

    runMonoM :: VNameSource -> MonoM a -> ((a, Seq.Seq (VName, ValBind)), VNameSource)
    runMonoM :: forall a.
    VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
    runMonoM VNameSource
    src (MonoM RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    m) = ((a
    a, Seq (VName, ValBind)
    defs), VNameSource
    src')
      where
        (a
    a, (ExpReplacements
    _, VNameSource
    src'), Seq (VName, ValBind)
    defs) = State
      Lifts (a, (ExpReplacements, VNameSource), Seq (VName, ValBind))
    -> Lifts
    -> (a, (ExpReplacements, VNameSource), Seq (VName, ValBind))
    forall s a. State s a -> s -> a
    evalState (RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> Env
    -> (ExpReplacements, VNameSource)
    -> State
         Lifts (a, (ExpReplacements, VNameSource), Seq (VName, ValBind))
    forall r w s (m :: * -> *) a.
    RWST r w s m a -> r -> s -> m (a, s, w)
    runRWST RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    m Env
    forall a. Monoid a => a
    mempty (ExpReplacements
    forall a. Monoid a => a
    mempty, VNameSource
    src)) Lifts
    forall a. Monoid a => a
    mempty

    lookupFun :: VName -> MonoM (Maybe PolyBinding)
    lookupFun :: VName -> MonoM (Maybe PolyBinding)
    lookupFun VName
    vn = do
      Map VName PolyBinding
    env <- (Env -> Map VName PolyBinding) -> MonoM (Map VName PolyBinding)
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks Env -> Map VName PolyBinding
    envPolyBindings
      case VName -> Map VName PolyBinding -> Maybe PolyBinding
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    vn Map VName PolyBinding
    env of
        Just PolyBinding
    valbind -> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Maybe PolyBinding -> MonoM (Maybe PolyBinding))
    -> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
    forall a b. (a -> b) -> a -> b
    $ PolyBinding -> Maybe PolyBinding
    forall a. a -> Maybe a
    Just PolyBinding
    valbind
        Maybe PolyBinding
    Nothing -> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe PolyBinding
    forall a. Maybe a
    Nothing

    askScope :: MonoM (S.Set VName)
    askScope :: MonoM (Set VName)
    askScope = do
      Set VName
    scope <- (Env -> Set VName) -> MonoM (Set VName)
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks Env -> Set VName
    envScope
      Set VName
    scope' <- (Env -> Set VName) -> MonoM (Set VName)
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks ((Env -> Set VName) -> MonoM (Set VName))
    -> (Env -> Set VName) -> MonoM (Set VName)
    forall a b. (a -> b) -> a -> b
    $ Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    S.union Set VName
    scope (Set VName -> Set VName) -> (Env -> Set VName) -> Env -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> Set VName
    envGlobalScope
      Set VName
    scope'' <- (Env -> Set VName) -> MonoM (Set VName)
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks ((Env -> Set VName) -> MonoM (Set VName))
    -> (Env -> Set VName) -> MonoM (Set VName)
    forall a b. (a -> b) -> a -> b
    $ Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    S.union Set VName
    scope' (Set VName -> Set VName) -> (Env -> Set VName) -> Env -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map VName PolyBinding -> Set VName
    forall k a. Map k a -> Set k
    M.keysSet (Map VName PolyBinding -> Set VName)
    -> (Env -> Map VName PolyBinding) -> Env -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> Map VName PolyBinding
    envPolyBindings
      Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    S.union Set VName
    scope'' (Set VName -> Set VName)
    -> (Lifts -> Set VName) -> Lifts -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList ([VName] -> Set VName) -> (Lifts -> [VName]) -> Lifts -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
    -> Lifts -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName, InferSizeArgs) -> VName
    forall a b. (a, b) -> a
    fst ((VName, InferSizeArgs) -> VName)
    -> (((VName, MonoType), (VName, InferSizeArgs))
        -> (VName, InferSizeArgs))
    -> ((VName, MonoType), (VName, InferSizeArgs))
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, MonoType), (VName, InferSizeArgs))
    -> (VName, InferSizeArgs)
    forall a b. (a, b) -> b
    snd) (Lifts -> Set VName) -> MonoM Lifts -> MonoM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM Lifts
    getLifts

    -- | Asks the introduced variables in a set of argument,
    -- that is arguments not currently in scope.
    askIntros :: S.Set VName -> MonoM (S.Set VName)
    askIntros :: Set VName -> MonoM (Set VName)
    askIntros Set VName
    argset =
      ((VName -> Bool) -> Set VName -> Set VName
    forall a. (a -> Bool) -> Set a -> Set a
    S.filter VName -> Bool
    notIntrisic Set VName
    argset `S.difference`) (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM (Set VName)
    askScope
      where
        notIntrisic :: VName -> Bool
    notIntrisic VName
    vn = VName -> Int
    baseTag VName
    vn Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    maxIntrinsicTag

    -- | Gets and removes expressions that could not be calculated when
    -- the arguments set will be unscoped.
    -- This should be called without argset in scope, for good detection of intros.
    parametrizing :: S.Set VName -> MonoM ExpReplacements
    parametrizing :: Set VName -> MonoM ExpReplacements
    parametrizing Set VName
    argset = do
      Set VName
    intros <- Set VName -> MonoM (Set VName)
    askIntros Set VName
    argset
      let usesIntros :: Exp -> Bool
    usesIntros = Bool -> Bool
    not (Bool -> Bool) -> (Exp -> Bool) -> Exp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Set VName -> Set VName -> Bool
    forall a. Ord a => Set a -> Set a -> Bool
    S.disjoint Set VName
    intros (Set VName -> Bool) -> (Exp -> Set VName) -> Exp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . FV -> Set VName
    fvVars (FV -> Set VName) -> (Exp -> FV) -> Exp -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Exp -> FV
    freeInExp
      (ExpReplacements
    params, ExpReplacements
    nxtBind) <- (ExpReplacements -> (ExpReplacements, ExpReplacements))
    -> MonoM (ExpReplacements, ExpReplacements)
    forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
    gets ((ExpReplacements -> (ExpReplacements, ExpReplacements))
     -> MonoM (ExpReplacements, ExpReplacements))
    -> (ExpReplacements -> (ExpReplacements, ExpReplacements))
    -> MonoM (ExpReplacements, ExpReplacements)
    forall a b. (a -> b) -> a -> b
    $ ((ReplacedExp, VName) -> Bool)
    -> ExpReplacements -> (ExpReplacements, ExpReplacements)
    forall a. (a -> Bool) -> [a] -> ([a], [a])
    partition (Exp -> Bool
    usesIntros (Exp -> Bool)
    -> ((ReplacedExp, VName) -> Exp) -> (ReplacedExp, VName) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ReplacedExp -> Exp
    unReplaced (ReplacedExp -> Exp)
    -> ((ReplacedExp, VName) -> ReplacedExp)
    -> (ReplacedExp, VName)
    -> Exp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (ReplacedExp, VName) -> ReplacedExp
    forall a b. (a, b) -> a
    fst)
      ExpReplacements -> MonoM ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put ExpReplacements
    nxtBind
      ExpReplacements -> MonoM ExpReplacements
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ExpReplacements
    params

    calculateDims :: Exp -> ExpReplacements -> MonoM Exp
    calculateDims :: Exp -> ExpReplacements -> MonoM Exp
    calculateDims Exp
    body ExpReplacements
    repl =
      ExpReplacements -> Exp -> MonoM Exp
    forall {f :: * -> *}.
    MonadFreshNames f =>
    ExpReplacements -> Exp -> f Exp
    foldCalc ExpReplacements
    top_repl (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ ExpReplacements -> Exp -> Exp
    expReplace ExpReplacements
    top_repl Exp
    body
      where
        depends :: (ReplacedExp, b) -> (ReplacedExp, b) -> Bool
    depends (ReplacedExp
    a, b
    _) (ReplacedExp
    b, b
    _) = ReplacedExp -> Exp
    unReplaced ReplacedExp
    b Exp -> [Exp] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` Exp -> [Exp]
    subExps (ReplacedExp -> Exp
    unReplaced ReplacedExp
    a)
        top_repl :: ExpReplacements
    top_repl = ((ReplacedExp, VName) -> (ReplacedExp, VName) -> Bool)
    -> ExpReplacements -> ExpReplacements
    forall a. (a -> a -> Bool) -> [a] -> [a]
    topologicalSort (ReplacedExp, VName) -> (ReplacedExp, VName) -> Bool
    forall {b} {b}. (ReplacedExp, b) -> (ReplacedExp, b) -> Bool
    depends ExpReplacements
    repl

        ---- Calculus insertion
        foldCalc :: ExpReplacements -> Exp -> f Exp
    foldCalc [] Exp
    body' = Exp -> f Exp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    body'
        foldCalc ((ReplacedExp
    dim, VName
    vn) : ExpReplacements
    repls) Exp
    body' = do
          VName
    reName <- VName -> f VName
    forall (m :: * -> *). MonadFreshNames m => VName -> m VName
    newName VName
    vn
          let expr :: Exp
    expr = ExpReplacements -> Exp -> Exp
    expReplace ExpReplacements
    repls (Exp -> Exp) -> Exp -> Exp
    forall a b. (a -> b) -> a -> b
    $ ReplacedExp -> Exp
    unReplaced ReplacedExp
    dim
              subst :: VName -> Maybe (Subst t)
    subst VName
    vn' =
                if VName
    vn' VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    vn
                  then Subst t -> Maybe (Subst t)
    forall a. a -> Maybe a
    Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
    forall a b. (a -> b) -> a -> b
    $ Exp -> Subst t
    forall t. Exp -> Subst t
    ExpSubst (Exp -> Subst t) -> Exp -> Subst t
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    reName) SrcLoc
    forall a. Monoid a => a
    mempty
                  else Maybe (Subst t)
    forall a. Maybe a
    Nothing
              appRes :: Info AppRes
    appRes = case Exp
    body' of
                (AppExp AppExpBase Info VName
    _ (Info (AppRes StructType
    ty [VName]
    ext))) -> 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 (TypeSubs -> StructType -> StructType
    forall a. Substitutable a => TypeSubs -> a -> a
    applySubst TypeSubs
    forall {t}. VName -> Maybe (Subst t)
    subst StructType
    ty) (VName
    reName VName -> [VName] -> [VName]
    forall a. a -> [a] -> [a]
    : [VName]
    ext)
                Exp
    e -> 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 (TypeSubs -> StructType -> StructType
    forall a. Substitutable a => TypeSubs -> a -> a
    applySubst TypeSubs
    forall {t}. VName -> Maybe (Subst t)
    subst (StructType -> StructType) -> StructType -> StructType
    forall a b. (a -> b) -> a -> b
    $ Exp -> StructType
    typeOf Exp
    e) [VName
    reName]
          ExpReplacements -> Exp -> f Exp
    foldCalc ExpReplacements
    repls (Exp -> f Exp) -> Exp -> f Exp
    forall a b. (a -> b) -> a -> b
    $
            AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp
              ( [SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    LetPat
                  []
                  (VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    vn (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    forall dim als. TypeBase dim als
    i64) (Exp -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Exp
    expr))
                  Exp
    expr
                  Exp
    body'
                  SrcLoc
    forall a. Monoid a => a
    mempty
              )
              Info AppRes
    appRes

    unscoping :: S.Set VName -> Exp -> MonoM Exp
    unscoping :: Set VName -> Exp -> MonoM Exp
    unscoping Set VName
    argset Exp
    body = do
      ExpReplacements
    localDims <- Set VName -> MonoM ExpReplacements
    parametrizing Set VName
    argset
      Set VName
    scope <- Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    S.union Set VName
    argset (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM (Set VName)
    askScope
      Exp -> ExpReplacements -> MonoM Exp
    calculateDims Exp
    body (ExpReplacements -> MonoM Exp) -> ExpReplacements -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Set VName -> ExpReplacements -> ExpReplacements
    canCalculate Set VName
    scope ExpReplacements
    localDims

    scoping :: S.Set VName -> MonoM Exp -> MonoM Exp
    scoping :: Set VName -> MonoM Exp -> MonoM Exp
    scoping Set VName
    argset MonoM Exp
    m =
      Set VName -> MonoM Exp -> MonoM Exp
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    argset MonoM Exp
    m MonoM Exp -> (Exp -> MonoM Exp) -> MonoM Exp
    forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Set VName -> Exp -> MonoM Exp
    unscoping Set VName
    argset

    -- Given instantiated type of function, produce size arguments.
    type InferSizeArgs = StructType -> MonoM [Exp]

    -- | The integer encodes an equivalence class, so we can keep
    -- track of sizes that are statically identical.
    data MonoSize
      = MonoKnown Int
      | MonoAnon Int
      deriving (MonoSize -> MonoSize -> Bool
    (MonoSize -> MonoSize -> Bool)
    -> (MonoSize -> MonoSize -> Bool) -> Eq MonoSize
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: MonoSize -> MonoSize -> Bool
    == :: MonoSize -> MonoSize -> Bool
    $c/= :: MonoSize -> MonoSize -> Bool
    /= :: MonoSize -> MonoSize -> Bool
    Eq, Int -> MonoSize -> ShowS
    [MonoSize] -> ShowS
    MonoSize -> [Char]
    (Int -> MonoSize -> ShowS)
    -> (MonoSize -> [Char]) -> ([MonoSize] -> ShowS) -> Show MonoSize
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> MonoSize -> ShowS
    showsPrec :: Int -> MonoSize -> ShowS
    $cshow :: MonoSize -> [Char]
    show :: MonoSize -> [Char]
    $cshowList :: [MonoSize] -> ShowS
    showList :: [MonoSize] -> ShowS
    Show)

    instance Pretty MonoSize where
      pretty :: forall ann. MonoSize -> Doc ann
    pretty (MonoKnown Int
    i) = Doc ann
    "?" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Int -> Doc ann
    forall ann. Int -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty Int
    i
      pretty (MonoAnon Int
    i) = Doc ann
    "??" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Int -> Doc ann
    forall ann. Int -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty Int
    i

    instance Pretty (Shape MonoSize) where
      pretty :: forall ann. Shape MonoSize -> Doc ann
    pretty (Shape [MonoSize]
    ds) = [Doc ann] -> Doc ann
    forall a. Monoid a => [a] -> a
    mconcat ((MonoSize -> Doc ann) -> [MonoSize] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    brackets (Doc ann -> Doc ann)
    -> (MonoSize -> Doc ann) -> MonoSize -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . MonoSize -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. MonoSize -> Doc ann
    pretty) [MonoSize]
    ds)

    -- The kind of type relative to which we monomorphise.  What is most
    -- important to us is not the specific dimensions, but merely whether
    -- they are known or anonymous/local.
    type MonoType = TypeBase MonoSize NoUniqueness

    monoType :: TypeBase Size als -> MonoType
    monoType :: forall als. TypeBase Exp als -> MonoType
    monoType = MonoType -> MonoType
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    noExts (MonoType -> MonoType)
    -> (TypeBase Exp als -> MonoType) -> TypeBase Exp als -> MonoType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (State (Int, Map Exp Int) MonoType -> (Int, Map Exp Int) -> MonoType
    forall s a. State s a -> s -> a
    `evalState` (Int
    0, Map Exp Int
    forall a. Monoid a => a
    mempty)) (State (Int, Map Exp Int) MonoType -> MonoType)
    -> (TypeBase Exp als -> State (Int, Map Exp Int) MonoType)
    -> TypeBase Exp als
    -> MonoType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Set VName
     -> DimPos -> Exp -> StateT (Int, Map Exp Int) Identity MonoSize)
    -> StructType -> State (Int, Map Exp Int) MonoType
    forall (f :: * -> *) fdim tdim als.
    Applicative f =>
    (Set VName -> DimPos -> fdim -> f tdim)
    -> TypeBase fdim als -> f (TypeBase tdim als)
    traverseDims Set VName
    -> DimPos -> Exp -> StateT (Int, Map Exp Int) Identity MonoSize
    forall {m :: * -> *} {p}.
    MonadState (Int, Map Exp Int) m =>
    Set VName -> p -> Exp -> m MonoSize
    onDim (StructType -> State (Int, Map Exp Int) MonoType)
    -> (TypeBase Exp als -> StructType)
    -> TypeBase Exp als
    -> State (Int, Map Exp Int) MonoType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase Exp als -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct
      where
        -- Remove exts from return types because we don't use them anymore.
        noExts :: TypeBase MonoSize u -> TypeBase MonoSize u
        noExts :: forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    noExts (Array u
    u Shape MonoSize
    shape ScalarTypeBase MonoSize NoUniqueness
    t) = u
    -> Shape MonoSize
    -> ScalarTypeBase MonoSize NoUniqueness
    -> TypeBase MonoSize u
    forall dim u.
    u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
    Array u
    u Shape MonoSize
    shape (ScalarTypeBase MonoSize NoUniqueness -> TypeBase MonoSize u)
    -> ScalarTypeBase MonoSize NoUniqueness -> TypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ ScalarTypeBase MonoSize NoUniqueness
    -> ScalarTypeBase MonoSize NoUniqueness
    forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    noExtsScalar ScalarTypeBase MonoSize NoUniqueness
    t
        noExts (Scalar ScalarTypeBase MonoSize u
    t) = ScalarTypeBase MonoSize u -> TypeBase MonoSize u
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase MonoSize u -> TypeBase MonoSize u)
    -> ScalarTypeBase MonoSize u -> TypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    noExtsScalar ScalarTypeBase MonoSize u
    t
        noExtsScalar :: ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    noExtsScalar (Record Map Name (TypeBase MonoSize u)
    fs) = Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    Record (Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u)
    -> Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ (TypeBase MonoSize u -> TypeBase MonoSize u)
    -> Map Name (TypeBase MonoSize u) -> Map Name (TypeBase MonoSize u)
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map TypeBase MonoSize u -> TypeBase MonoSize u
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    noExts Map Name (TypeBase MonoSize u)
    fs
        noExtsScalar (Sum Map Name [TypeBase MonoSize u]
    fs) = Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
    forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
    Sum (Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u)
    -> Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ ([TypeBase MonoSize u] -> [TypeBase MonoSize u])
    -> Map Name [TypeBase MonoSize u] -> Map Name [TypeBase MonoSize u]
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map ((TypeBase MonoSize u -> TypeBase MonoSize u)
    -> [TypeBase MonoSize u] -> [TypeBase MonoSize u]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase MonoSize u -> TypeBase MonoSize u
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    noExts) Map Name [TypeBase MonoSize u]
    fs
        noExtsScalar (Arrow u
    as PName
    p Diet
    d MonoType
    t1 (RetType [VName]
    _ TypeBase MonoSize Uniqueness
    t2)) =
          u
    -> PName
    -> Diet
    -> MonoType
    -> RetTypeBase MonoSize Uniqueness
    -> ScalarTypeBase MonoSize u
    forall dim u.
    u
    -> PName
    -> Diet
    -> TypeBase dim NoUniqueness
    -> RetTypeBase dim Uniqueness
    -> ScalarTypeBase dim u
    Arrow u
    as PName
    p Diet
    d (MonoType -> MonoType
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    noExts MonoType
    t1) ([VName]
    -> TypeBase MonoSize Uniqueness -> RetTypeBase MonoSize Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [] (TypeBase MonoSize Uniqueness -> TypeBase MonoSize Uniqueness
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    noExts TypeBase MonoSize Uniqueness
    t2))
        noExtsScalar ScalarTypeBase MonoSize u
    t = ScalarTypeBase MonoSize u
    t
        onDim :: Set VName -> p -> Exp -> m MonoSize
    onDim Set VName
    bound p
    _ Exp
    d
          -- A locally bound size.
          | (VName -> Bool) -> Set VName -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.member` Set VName
    bound) (Set VName -> Bool) -> Set VName -> Bool
    forall a b. (a -> b) -> a -> b
    $ FV -> Set VName
    fvVars (FV -> Set VName) -> FV -> Set VName
    forall a b. (a -> b) -> a -> b
    $ Exp -> FV
    freeInExp Exp
    d = do
              (Int
    i, Map Exp Int
    m) <- m (Int, Map Exp Int)
    forall s (m :: * -> *). MonadState s m => m s
    get
              case Exp -> Map Exp Int -> Maybe Int
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Exp
    d Map Exp Int
    m of
                Just Int
    prev ->
                  MonoSize -> m MonoSize
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
    forall a b. (a -> b) -> a -> b
    $ Int -> MonoSize
    MonoAnon Int
    prev
                Maybe Int
    Nothing -> do
                  (Int, Map Exp Int) -> m ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put (Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1, Exp -> Int -> Map Exp Int -> Map Exp Int
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert Exp
    d Int
    i Map Exp Int
    m)
                  MonoSize -> m MonoSize
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
    forall a b. (a -> b) -> a -> b
    $ Int -> MonoSize
    MonoAnon Int
    i
        onDim Set VName
    _ p
    _ Exp
    d = do
          (Int
    i, Map Exp Int
    m) <- m (Int, Map Exp Int)
    forall s (m :: * -> *). MonadState s m => m s
    get
          case Exp -> Map Exp Int -> Maybe Int
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Exp
    d Map Exp Int
    m of
            Just Int
    prev ->
              MonoSize -> m MonoSize
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
    forall a b. (a -> b) -> a -> b
    $ Int -> MonoSize
    MonoKnown Int
    prev
            Maybe Int
    Nothing -> do
              -- Ensure that each instance of anySize is treated distinctly.
              (Int, Map Exp Int) -> m ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put (Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1, if Exp
    d Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    == Exp
    anySize then Map Exp Int
    m else Exp -> Int -> Map Exp Int -> Map Exp Int
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert Exp
    d Int
    i Map Exp Int
    m)
              MonoSize -> m MonoSize
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
    forall a b. (a -> b) -> a -> b
    $ Int -> MonoSize
    MonoKnown Int
    i

    -- Mapping from function name and instance list to a new function name in case
    -- the function has already been instantiated with those concrete types.
    type Lifts = [((VName, MonoType), (VName, InferSizeArgs))]

    getLifts :: MonoM Lifts
    getLifts :: MonoM Lifts
    getLifts = RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      Lifts
    -> MonoM Lifts
    forall a.
    RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> MonoM a
    MonoM (RWST
       Env
       (Seq (VName, ValBind))
       (ExpReplacements, VNameSource)
       (State Lifts)
       Lifts
     -> MonoM Lifts)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         Lifts
    -> MonoM Lifts
    forall a b. (a -> b) -> a -> b
    $ State Lifts Lifts
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         Lifts
    forall (m :: * -> *) a.
    Monad m =>
    m a
    -> RWST
         Env (Seq (VName, ValBind)) (ExpReplacements, VNameSource) m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift State Lifts Lifts
    forall s (m :: * -> *). MonadState s m => m s
    get

    modifyLifts :: (Lifts -> Lifts) -> MonoM ()
    modifyLifts :: (Lifts -> Lifts) -> MonoM ()
    modifyLifts = RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ()
    -> MonoM ()
    forall a.
    RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      a
    -> MonoM a
    MonoM (RWST
       Env
       (Seq (VName, ValBind))
       (ExpReplacements, VNameSource)
       (State Lifts)
       ()
     -> MonoM ())
    -> ((Lifts -> Lifts)
        -> RWST
             Env
             (Seq (VName, ValBind))
             (ExpReplacements, VNameSource)
             (State Lifts)
             ())
    -> (Lifts -> Lifts)
    -> MonoM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . State Lifts ()
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ()
    forall (m :: * -> *) a.
    Monad m =>
    m a
    -> RWST
         Env (Seq (VName, ValBind)) (ExpReplacements, VNameSource) m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (State Lifts ()
     -> RWST
          Env
          (Seq (VName, ValBind))
          (ExpReplacements, VNameSource)
          (State Lifts)
          ())
    -> ((Lifts -> Lifts) -> State Lifts ())
    -> (Lifts -> Lifts)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Lifts -> Lifts) -> State Lifts ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify

    addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
    addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
    addLifted VName
    fname MonoType
    il (VName, InferSizeArgs)
    liftf =
      (Lifts -> Lifts) -> MonoM ()
    modifyLifts (((VName
    fname, MonoType
    il), (VName, InferSizeArgs)
    liftf) :)

    lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
    lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
    lookupLifted VName
    fname MonoType
    t = (VName, MonoType) -> Lifts -> Maybe (VName, InferSizeArgs)
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup (VName
    fname, MonoType
    t) (Lifts -> Maybe (VName, InferSizeArgs))
    -> MonoM Lifts -> MonoM (Maybe (VName, InferSizeArgs))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM Lifts
    getLifts

    sizeVarName :: Exp -> String
    sizeVarName :: Exp -> [Char]
    sizeVarName Exp
    e = [Char]
    "d<{" [Char] -> ShowS
    forall a. Semigroup a => a -> a -> a
    <> ExpBase NoInfo VName -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString (Exp -> ExpBase NoInfo VName
    bareExp Exp
    e) [Char] -> ShowS
    forall a. Semigroup a => a -> a -> a
    <> [Char]
    "}>"

    -- | Creates a new expression replacement if needed, this always produces normalised sizes.
    -- (e.g. single variable or constant)
    replaceExp :: Exp -> MonoM Exp
    replaceExp :: Exp -> MonoM Exp
    replaceExp Exp
    e =
      case Exp -> Maybe Exp
    maybeNormalisedSize Exp
    e of
        Just Exp
    e' -> Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e'
        Maybe Exp
    Nothing -> do
          let e' :: ReplacedExp
    e' = Exp -> ReplacedExp
    ReplacedExp Exp
    e
          Maybe VName
    prev <- (ExpReplacements -> Maybe VName) -> MonoM (Maybe VName)
    forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
    gets ((ExpReplacements -> Maybe VName) -> MonoM (Maybe VName))
    -> (ExpReplacements -> Maybe VName) -> MonoM (Maybe VName)
    forall a b. (a -> b) -> a -> b
    $ ReplacedExp -> ExpReplacements -> Maybe VName
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup ReplacedExp
    e'
          Maybe VName
    prev_param <- (Env -> Maybe VName) -> MonoM (Maybe VName)
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks ((Env -> Maybe VName) -> MonoM (Maybe VName))
    -> (Env -> Maybe VName) -> MonoM (Maybe VName)
    forall a b. (a -> b) -> a -> b
    $ ReplacedExp -> ExpReplacements -> Maybe VName
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup ReplacedExp
    e' (ExpReplacements -> Maybe VName)
    -> (Env -> ExpReplacements) -> Env -> Maybe VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> ExpReplacements
    envParametrized
          case (Maybe VName
    prev_param, Maybe VName
    prev) of
            (Just VName
    vn, Maybe VName
    _) -> Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    vn) (Exp -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Exp
    e)
            (Maybe VName
    Nothing, Just VName
    vn) -> Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    vn) (Exp -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Exp
    e)
            (Maybe VName
    Nothing, Maybe VName
    Nothing) -> do
              VName
    vn <- [Char] -> MonoM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newNameFromString ([Char] -> MonoM VName) -> [Char] -> MonoM VName
    forall a b. (a -> b) -> a -> b
    $ Exp -> [Char]
    sizeVarName Exp
    e
              (ExpReplacements -> ExpReplacements) -> MonoM ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((ReplacedExp
    e', VName
    vn) :)
              Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    vn) (Exp -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Exp
    e)
      where
        -- Avoid replacing of some 'already normalised' sizes that are just surounded by some parentheses.
        maybeNormalisedSize :: Exp -> Maybe Exp
    maybeNormalisedSize Exp
    e'
          | Just Exp
    e'' <- Exp -> Maybe Exp
    stripExp Exp
    e' = Exp -> Maybe Exp
    maybeNormalisedSize Exp
    e''
        maybeNormalisedSize (Var QualName VName
    qn Info StructType
    _ SrcLoc
    loc) = Exp -> Maybe Exp
    forall a. a -> Maybe a
    Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName QualName VName
    qn SrcLoc
    loc
        maybeNormalisedSize (IntLit Integer
    v Info StructType
    _ SrcLoc
    loc) = Exp -> Maybe Exp
    forall a. a -> Maybe a
    Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
    forall a b. (a -> b) -> a -> b
    $ Integer -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    Integer -> f StructType -> SrcLoc -> ExpBase f vn
    IntLit Integer
    v (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    forall dim als. TypeBase dim als
    i64) SrcLoc
    loc
        maybeNormalisedSize Exp
    _ = Maybe Exp
    forall a. Maybe a
    Nothing

    transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
    transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
    transformFName SrcLoc
    loc QualName VName
    fname StructType
    ft = do
      StructType
    t' <- StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    ft
      let mono_t :: MonoType
    mono_t = StructType -> MonoType
    forall als. TypeBase Exp als -> MonoType
    monoType StructType
    ft
      if VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    fname) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag
        then Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> StructType -> Exp
    forall {vn}. QualName vn -> StructType -> ExpBase Info vn
    var QualName VName
    fname StructType
    t'
        else do
          Maybe (VName, InferSizeArgs)
    maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
    lookupLifted (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    fname) MonoType
    mono_t
          Maybe PolyBinding
    maybe_funbind <- VName -> MonoM (Maybe PolyBinding)
    lookupFun (VName -> MonoM (Maybe PolyBinding))
    -> VName -> MonoM (Maybe PolyBinding)
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    fname
          case (Maybe (VName, InferSizeArgs)
    maybe_fname, Maybe PolyBinding
    maybe_funbind) of
            -- The function has already been monomorphised.
            (Just (VName
    fname', InferSizeArgs
    infer), Maybe PolyBinding
    _) ->
              VName -> TypeBase Exp Uniqueness -> [Exp] -> Exp
    forall {vn}.
    vn
    -> TypeBase Exp Uniqueness -> [ExpBase Info vn] -> ExpBase Info vn
    applySizeArgs VName
    fname' (Uniqueness -> StructType -> TypeBase Exp Uniqueness
    forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
    toRes Uniqueness
    Nonunique StructType
    t') ([Exp] -> Exp) -> MonoM [Exp] -> MonoM Exp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> InferSizeArgs
    infer StructType
    t'
            -- An intrinsic function.
            (Maybe (VName, InferSizeArgs)
    Nothing, Maybe PolyBinding
    Nothing) -> Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> StructType -> Exp
    forall {vn}. QualName vn -> StructType -> ExpBase Info vn
    var QualName VName
    fname StructType
    t'
            -- A polymorphic function.
            (Maybe (VName, InferSizeArgs)
    Nothing, Just PolyBinding
    funbind) -> do
              (VName
    fname', InferSizeArgs
    infer, ValBind
    funbind') <- PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
    monomorphiseBinding PolyBinding
    funbind MonoType
    mono_t
              Seq (VName, ValBind) -> MonoM ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell (Seq (VName, ValBind) -> MonoM ())
    -> Seq (VName, ValBind) -> MonoM ()
    forall a b. (a -> b) -> a -> b
    $ (VName, ValBind) -> Seq (VName, ValBind)
    forall a. a -> Seq a
    Seq.singleton (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    fname, ValBind
    funbind')
              VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
    addLifted (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    fname) MonoType
    mono_t (VName
    fname', InferSizeArgs
    infer)
              VName -> TypeBase Exp Uniqueness -> [Exp] -> Exp
    forall {vn}.
    vn
    -> TypeBase Exp Uniqueness -> [ExpBase Info vn] -> ExpBase Info vn
    applySizeArgs VName
    fname' (Uniqueness -> StructType -> TypeBase Exp Uniqueness
    forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
    toRes Uniqueness
    Nonunique StructType
    t') ([Exp] -> Exp) -> MonoM [Exp] -> MonoM Exp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> InferSizeArgs
    infer StructType
    t'
      where
        var :: QualName vn -> StructType -> ExpBase Info vn
    var QualName vn
    fname' StructType
    t' = QualName vn -> Info StructType -> SrcLoc -> ExpBase Info vn
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var QualName vn
    fname' (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    t') SrcLoc
    loc

        applySizeArg :: TypeBase Exp Uniqueness
    -> (Int, ExpBase Info vn)
    -> ExpBase Info vn
    -> (Int, ExpBase Info vn)
    applySizeArg TypeBase Exp Uniqueness
    t (Int
    i, ExpBase Info vn
    f) ExpBase Info vn
    size_arg =
          ( Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1,
            ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply
              ExpBase Info vn
    f
              [(Maybe VName
    forall a. Maybe a
    Nothing, ExpBase Info vn
    size_arg)]
              (StructType -> [VName] -> AppRes
    AppRes ([ParamType] -> RetTypeBase Exp Uniqueness -> StructType
    foldFunType (Int -> ParamType -> [ParamType]
    forall a. Int -> a -> [a]
    replicate Int
    i ParamType
    forall dim als. TypeBase dim als
    i64) ([VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [] TypeBase Exp Uniqueness
    t)) [])
          )

        applySizeArgs :: vn
    -> TypeBase Exp Uniqueness -> [ExpBase Info vn] -> ExpBase Info vn
    applySizeArgs vn
    fname' TypeBase Exp Uniqueness
    t [ExpBase Info vn]
    size_args =
          (Int, ExpBase Info vn) -> ExpBase Info vn
    forall a b. (a, b) -> b
    snd ((Int, ExpBase Info vn) -> ExpBase Info vn)
    -> (Int, ExpBase Info vn) -> ExpBase Info vn
    forall a b. (a -> b) -> a -> b
    $
            ((Int, ExpBase Info vn)
     -> ExpBase Info vn -> (Int, ExpBase Info vn))
    -> (Int, ExpBase Info vn)
    -> [ExpBase Info vn]
    -> (Int, ExpBase Info vn)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl'
              (TypeBase Exp Uniqueness
    -> (Int, ExpBase Info vn)
    -> ExpBase Info vn
    -> (Int, ExpBase Info vn)
    forall {vn}.
    TypeBase Exp Uniqueness
    -> (Int, ExpBase Info vn)
    -> ExpBase Info vn
    -> (Int, ExpBase Info vn)
    applySizeArg TypeBase Exp Uniqueness
    t)
              ( [ExpBase Info vn] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [ExpBase Info vn]
    size_args Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1,
                QualName vn -> Info StructType -> SrcLoc -> ExpBase Info vn
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var
                  (vn -> QualName vn
    forall v. v -> QualName v
    qualName vn
    fname')
                  (StructType -> Info StructType
    forall a. a -> Info a
    Info ([ParamType] -> RetTypeBase Exp Uniqueness -> StructType
    foldFunType ((ExpBase Info vn -> ParamType) -> [ExpBase Info vn] -> [ParamType]
    forall a b. (a -> b) -> [a] -> [b]
    map (ParamType -> ExpBase Info vn -> ParamType
    forall a b. a -> b -> a
    const ParamType
    forall dim als. TypeBase dim als
    i64) [ExpBase Info vn]
    size_args) ([VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [] TypeBase Exp Uniqueness
    t)))
                  SrcLoc
    loc
              )
              [ExpBase Info vn]
    size_args

    transformType :: TypeBase Size u -> MonoM (TypeBase Size u)
    transformType :: forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType TypeBase Exp u
    typ =
      case TypeBase Exp u
    typ of
        Scalar ScalarTypeBase Exp u
    scalar ->
          ScalarTypeBase Exp u -> TypeBase Exp u
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase Exp u -> TypeBase Exp u)
    -> MonoM (ScalarTypeBase Exp u) -> MonoM (TypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
    forall u. ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
    transformScalarSizes ScalarTypeBase Exp u
    scalar
        Array u
    u Shape Exp
    shape ScalarTypeBase Exp NoUniqueness
    scalar ->
          u -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u
    forall dim u.
    u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
    Array u
    u (Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u)
    -> MonoM (Shape Exp)
    -> MonoM (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> MonoM Exp) -> Shape Exp -> MonoM (Shape 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) -> Shape a -> m (Shape b)
    mapM Exp -> MonoM Exp
    onDim Shape Exp
    shape MonoM (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u)
    -> MonoM (ScalarTypeBase Exp NoUniqueness)
    -> MonoM (TypeBase Exp u)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> ScalarTypeBase Exp NoUniqueness
    -> MonoM (ScalarTypeBase Exp NoUniqueness)
    forall u. ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
    transformScalarSizes ScalarTypeBase Exp NoUniqueness
    scalar
      where
        transformScalarSizes :: ScalarTypeBase Size u -> MonoM (ScalarTypeBase Size u)
        transformScalarSizes :: forall u. ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
    transformScalarSizes (Record Map Name (TypeBase Exp u)
    fs) =
          Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    Record (Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u)
    -> MonoM (Map Name (TypeBase Exp u))
    -> MonoM (ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (TypeBase Exp u -> MonoM (TypeBase Exp u))
    -> Map Name (TypeBase Exp u) -> MonoM (Map Name (TypeBase Exp u))
    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) -> Map Name a -> f (Map Name b)
    traverse TypeBase Exp u -> MonoM (TypeBase Exp u)
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Map Name (TypeBase Exp u)
    fs
        transformScalarSizes (Sum Map Name [TypeBase Exp u]
    cs) =
          Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u
    forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
    Sum (Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u)
    -> MonoM (Map Name [TypeBase Exp u])
    -> MonoM (ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (([TypeBase Exp u] -> MonoM [TypeBase Exp u])
    -> Map Name [TypeBase Exp u] -> MonoM (Map Name [TypeBase Exp u])
    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) -> Map Name a -> f (Map Name b)
    traverse (([TypeBase Exp u] -> MonoM [TypeBase Exp u])
     -> Map Name [TypeBase Exp u] -> MonoM (Map Name [TypeBase Exp u]))
    -> ((TypeBase Exp u -> MonoM (TypeBase Exp u))
        -> [TypeBase Exp u] -> MonoM [TypeBase Exp u])
    -> (TypeBase Exp u -> MonoM (TypeBase Exp u))
    -> Map Name [TypeBase Exp u]
    -> MonoM (Map Name [TypeBase Exp u])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (TypeBase Exp u -> MonoM (TypeBase Exp u))
    -> [TypeBase Exp u] -> MonoM [TypeBase Exp u]
    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) TypeBase Exp u -> MonoM (TypeBase Exp u)
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Map Name [TypeBase Exp u]
    cs
        transformScalarSizes (Arrow u
    as PName
    argName Diet
    d StructType
    argT RetTypeBase Exp Uniqueness
    retT) =
          u
    -> PName
    -> Diet
    -> StructType
    -> RetTypeBase Exp Uniqueness
    -> ScalarTypeBase Exp u
    forall dim u.
    u
    -> PName
    -> Diet
    -> TypeBase dim NoUniqueness
    -> RetTypeBase dim Uniqueness
    -> ScalarTypeBase dim u
    Arrow u
    as PName
    argName Diet
    d
            (StructType -> RetTypeBase Exp Uniqueness -> ScalarTypeBase Exp u)
    -> MonoM StructType
    -> MonoM (RetTypeBase Exp Uniqueness -> ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    argT
            MonoM (RetTypeBase Exp Uniqueness -> ScalarTypeBase Exp u)
    -> MonoM (RetTypeBase Exp Uniqueness)
    -> MonoM (ScalarTypeBase Exp u)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Set VName
    -> RetTypeBase Exp Uniqueness -> MonoM (RetTypeBase Exp Uniqueness)
    forall as.
    Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    transformRetTypeSizes Set VName
    argset RetTypeBase Exp Uniqueness
    retT
          where
            argset :: Set VName
    argset =
              case PName
    argName of
                PName
    Unnamed -> Set VName
    forall a. Monoid a => a
    mempty
                Named VName
    vn -> VName -> Set VName
    forall a. a -> Set a
    S.singleton VName
    vn
        transformScalarSizes (TypeVar u
    u QualName VName
    qn [TypeArg Exp]
    args) =
          u -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp u
    forall dim u.
    u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
    TypeVar u
    u QualName VName
    qn ([TypeArg Exp] -> ScalarTypeBase Exp u)
    -> MonoM [TypeArg Exp] -> MonoM (ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (TypeArg Exp -> MonoM (TypeArg Exp))
    -> [TypeArg Exp] -> MonoM [TypeArg 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 TypeArg Exp -> MonoM (TypeArg Exp)
    onArg [TypeArg Exp]
    args
          where
            onArg :: TypeArg Exp -> MonoM (TypeArg Exp)
    onArg (TypeArgDim Exp
    dim) = Exp -> TypeArg Exp
    forall dim. dim -> TypeArg dim
    TypeArgDim (Exp -> TypeArg Exp) -> MonoM Exp -> MonoM (TypeArg Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    onDim Exp
    dim
            onArg (TypeArgType StructType
    ty) = StructType -> TypeArg Exp
    forall dim. TypeBase dim NoUniqueness -> TypeArg dim
    TypeArgType (StructType -> TypeArg Exp)
    -> MonoM StructType -> MonoM (TypeArg Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    ty
        transformScalarSizes ty :: ScalarTypeBase Exp u
    ty@Prim {} = ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ScalarTypeBase Exp u
    ty

        onDim :: Exp -> MonoM Exp
    onDim Exp
    e
          | Exp
    e Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    == Exp
    anySize = Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e
          | Bool
    otherwise = Exp -> MonoM Exp
    replaceExp (Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Exp -> MonoM Exp
    transformExp Exp
    e

    transformRetTypeSizes :: S.Set VName -> RetTypeBase Size as -> MonoM (RetTypeBase Size as)
    transformRetTypeSizes :: forall as.
    Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    transformRetTypeSizes Set VName
    argset (RetType [VName]
    dims TypeBase Exp as
    ty) = do
      TypeBase Exp as
    ty' <- Set VName -> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    argset (MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as))
    -> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
    forall a b. (a -> b) -> a -> b
    $ [VName] -> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
    forall a. [VName] -> MonoM a -> MonoM a
    withMono [VName]
    dims (MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as))
    -> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
    forall a b. (a -> b) -> a -> b
    $ TypeBase Exp as -> MonoM (TypeBase Exp as)
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType TypeBase Exp as
    ty
      ExpReplacements
    rl <- Set VName -> MonoM ExpReplacements
    parametrizing Set VName
    argset
      let dims' :: [VName]
    dims' = [VName]
    dims [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> ((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (ReplacedExp, VName) -> VName
    forall a b. (a, b) -> b
    snd ExpReplacements
    rl
      RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (RetTypeBase Exp as -> MonoM (RetTypeBase Exp as))
    -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    forall a b. (a -> b) -> a -> b
    $ [VName] -> TypeBase Exp as -> RetTypeBase Exp as
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [VName]
    dims' TypeBase Exp as
    ty'

    sizesForPat :: (MonadFreshNames m) => Pat ParamType -> m ([VName], Pat ParamType)
    sizesForPat :: forall (m :: * -> *).
    MonadFreshNames m =>
    Pat ParamType -> m ([VName], Pat ParamType)
    sizesForPat Pat ParamType
    pat = do
      (Pat ParamType
    params', [VName]
    sizes) <- StateT [VName] m (Pat ParamType)
    -> [VName] -> m (Pat ParamType, [VName])
    forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
    runStateT ((ParamType -> StateT [VName] m ParamType)
    -> Pat ParamType -> StateT [VName] m (Pat ParamType)
    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) -> PatBase Info VName a -> f (PatBase Info VName b)
    traverse ((Exp -> StateT [VName] m Exp)
    -> (Diet -> StateT [VName] m Diet)
    -> ParamType
    -> StateT [VName] m ParamType
    forall (f :: * -> *) a c b d.
    Applicative f =>
    (a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
    forall (t :: * -> * -> *) (f :: * -> *) a c b d.
    (Bitraversable t, Applicative f) =>
    (a -> f c) -> (b -> f d) -> t a b -> f (t c d)
    bitraverse Exp -> StateT [VName] m Exp
    forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
    (MonadTrans t, MonadFreshNames m, MonadState [VName] (t m)) =>
    Exp -> t m Exp
    onDim Diet -> StateT [VName] m Diet
    forall a. a -> StateT [VName] m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure) Pat ParamType
    pat) []
      ([VName], Pat ParamType) -> m ([VName], Pat ParamType)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([VName]
    sizes, Pat ParamType
    params')
      where
        onDim :: Exp -> t m Exp
    onDim Exp
    d
          | Exp
    d Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    == Exp
    anySize = do
              VName
    v <- m VName -> t m VName
    forall (m :: * -> *) a. Monad m => m a -> t m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m VName -> t m VName) -> m VName -> t m VName
    forall a b. (a -> b) -> a -> b
    $ [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "size"
              ([VName] -> [VName]) -> t m ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify (VName
    v :)
              Exp -> t m Exp
    forall a. a -> t m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> t m Exp) -> Exp -> t m Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    v) SrcLoc
    forall a. Monoid a => a
    mempty
          | Bool
    otherwise = Exp -> t m Exp
    forall a. a -> t m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    d

    transformAppRes :: AppRes -> MonoM AppRes
    transformAppRes :: AppRes -> MonoM AppRes
    transformAppRes (AppRes StructType
    t [VName]
    ext) =
      StructType -> [VName] -> AppRes
    AppRes (StructType -> [VName] -> AppRes)
    -> MonoM StructType -> MonoM ([VName] -> AppRes)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    t MonoM ([VName] -> AppRes) -> MonoM [VName] -> MonoM AppRes
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [VName] -> MonoM [VName]
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [VName]
    ext

    transformAppExp :: AppExp -> AppRes -> MonoM Exp
    transformAppExp :: AppExpBase Info VName -> AppRes -> MonoM Exp
    transformAppExp (Range Exp
    e1 Maybe Exp
    me Inclusiveness Exp
    incl SrcLoc
    loc) AppRes
    res = do
      Exp
    e1' <- Exp -> MonoM Exp
    transformExp Exp
    e1
      Maybe Exp
    me' <- (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe 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) -> Maybe a -> m (Maybe b)
    mapM Exp -> MonoM Exp
    transformExp Maybe Exp
    me
      Inclusiveness Exp
    incl' <- (Exp -> MonoM Exp)
    -> Inclusiveness Exp -> MonoM (Inclusiveness 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) -> Inclusiveness a -> m (Inclusiveness b)
    mapM Exp -> MonoM Exp
    transformExp Inclusiveness Exp
    incl
      AppRes
    res' <- AppRes -> MonoM AppRes
    transformAppRes AppRes
    res
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp (Exp
    -> Maybe Exp
    -> Inclusiveness Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> Maybe (ExpBase f vn)
    -> Inclusiveness (ExpBase f vn)
    -> SrcLoc
    -> AppExpBase f vn
    Range Exp
    e1' Maybe Exp
    me' Inclusiveness Exp
    incl' SrcLoc
    loc) (AppRes -> Info AppRes
    forall a. a -> Info a
    Info AppRes
    res')
    transformAppExp (LetPat [SizeBinder VName]
    sizes PatBase Info VName StructType
    pat Exp
    e Exp
    body SrcLoc
    loc) AppRes
    res = do
      Exp
    e' <- Exp -> MonoM Exp
    transformExp Exp
    e
      let dimArgs :: Set VName
    dimArgs = [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList ((SizeBinder VName -> VName) -> [SizeBinder VName] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SizeBinder VName -> VName
    forall vn. SizeBinder vn -> vn
    sizeName [SizeBinder VName]
    sizes)
      Set VName
    implicitDims <- Set VName -> MonoM (Set VName) -> MonoM (Set VName)
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    dimArgs (MonoM (Set VName) -> MonoM (Set VName))
    -> MonoM (Set VName) -> MonoM (Set VName)
    forall a b. (a -> b) -> a -> b
    $ Set VName -> MonoM (Set VName)
    askIntros (Set VName -> MonoM (Set VName)) -> Set VName -> MonoM (Set VName)
    forall a b. (a -> b) -> a -> b
    $ FV -> Set VName
    fvVars (FV -> Set VName) -> FV -> Set VName
    forall a b. (a -> b) -> a -> b
    $ PatBase Info VName StructType -> FV
    forall u. Pat (TypeBase Exp u) -> FV
    freeInPat PatBase Info VName StructType
    pat
      let dimArgs' :: Set VName
    dimArgs' = Set VName
    dimArgs Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> Set VName
    implicitDims
          letArgs :: Set VName
    letArgs = [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
    forall a b. (a -> b) -> a -> b
    $ PatBase Info VName StructType -> [VName]
    forall t. Pat t -> [VName]
    patNames PatBase Info VName StructType
    pat
          argset :: Set VName
    argset = Set VName
    dimArgs' Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` Set VName
    letArgs
      PatBase Info VName StructType
    pat' <- Set VName
    -> MonoM (PatBase Info VName StructType)
    -> MonoM (PatBase Info VName StructType)
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    dimArgs' (MonoM (PatBase Info VName StructType)
     -> MonoM (PatBase Info VName StructType))
    -> MonoM (PatBase Info VName StructType)
    -> MonoM (PatBase Info VName StructType)
    forall a b. (a -> b) -> a -> b
    $ PatBase Info VName StructType
    -> MonoM (PatBase Info VName StructType)
    forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
    transformPat PatBase Info VName StructType
    pat
      ExpReplacements
    params <- Set VName -> MonoM ExpReplacements
    parametrizing Set VName
    dimArgs'
      let sizes' :: [SizeBinder VName]
    sizes' = [SizeBinder VName]
    sizes [SizeBinder VName] -> [SizeBinder VName] -> [SizeBinder VName]
    forall a. Semigroup a => a -> a -> a
    <> (VName -> SizeBinder VName) -> [VName] -> [SizeBinder VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> SrcLoc -> SizeBinder VName
    forall vn. vn -> SrcLoc -> SizeBinder vn
    `SizeBinder` SrcLoc
    forall a. Monoid a => a
    mempty) (((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (ReplacedExp, VName) -> VName
    forall a b. (a, b) -> b
    snd ExpReplacements
    params [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> Set VName -> [VName]
    forall a. Set a -> [a]
    S.toList Set VName
    implicitDims)
      Exp
    body' <- ExpReplacements -> MonoM Exp -> MonoM Exp
    forall a. ExpReplacements -> MonoM a -> MonoM a
    withParams ExpReplacements
    params (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Set VName -> MonoM Exp -> MonoM Exp
    scoping Set VName
    argset (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Exp -> MonoM Exp
    transformExp Exp
    body
      AppRes
    res' <- AppRes -> MonoM AppRes
    transformAppRes AppRes
    res
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp ([SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    LetPat [SizeBinder VName]
    sizes' PatBase Info VName StructType
    pat' Exp
    e' Exp
    body' SrcLoc
    loc) (AppRes -> Info AppRes
    forall a. a -> Info a
    Info AppRes
    res')
    transformAppExp LetFun {} AppRes
    _ =
      [Char] -> MonoM Exp
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "transformAppExp: LetFun is not supposed to occur"
    transformAppExp (If Exp
    e1 Exp
    e2 Exp
    e3 SrcLoc
    loc) AppRes
    res =
      AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp (AppExpBase Info VName -> Info AppRes -> Exp)
    -> MonoM (AppExpBase Info VName) -> MonoM (Info AppRes -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
    If (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName)
    -> MonoM Exp
    -> MonoM (Exp -> Exp -> SrcLoc -> AppExpBase Info VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e1 MonoM (Exp -> Exp -> SrcLoc -> AppExpBase Info VName)
    -> MonoM Exp -> MonoM (Exp -> SrcLoc -> AppExpBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Exp -> MonoM Exp
    transformExp Exp
    e2 MonoM (Exp -> SrcLoc -> AppExpBase Info VName)
    -> MonoM Exp -> MonoM (SrcLoc -> AppExpBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Exp -> MonoM Exp
    transformExp Exp
    e3 MonoM (SrcLoc -> AppExpBase Info VName)
    -> MonoM SrcLoc -> MonoM (AppExpBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (AppRes -> Info AppRes
    forall a. a -> Info a
    Info (AppRes -> Info AppRes) -> MonoM AppRes -> MonoM (Info AppRes)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> AppRes -> MonoM AppRes
    transformAppRes AppRes
    res)
    transformAppExp (Apply Exp
    fe NonEmpty (Info (Maybe VName), Exp)
    args SrcLoc
    _) AppRes
    res =
      Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply
        (Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp)
    -> MonoM Exp -> MonoM ([(Maybe VName, Exp)] -> AppRes -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    fe
        MonoM ([(Maybe VName, Exp)] -> AppRes -> Exp)
    -> MonoM [(Maybe VName, Exp)] -> MonoM (AppRes -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> ((Info (Maybe VName), Exp) -> MonoM (Maybe VName, Exp))
    -> [(Info (Maybe VName), Exp)] -> MonoM [(Maybe VName, 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 (Info (Maybe VName), Exp) -> MonoM (Maybe VName, Exp)
    forall {t}. (Info t, Exp) -> MonoM (t, Exp)
    onArg (NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
    forall a. NonEmpty a -> [a]
    NE.toList NonEmpty (Info (Maybe VName), Exp)
    args)
        MonoM (AppRes -> Exp) -> MonoM AppRes -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> AppRes -> MonoM AppRes
    transformAppRes AppRes
    res
      where
        onArg :: (Info t, Exp) -> MonoM (t, Exp)
    onArg (Info t
    ext, Exp
    e) = (t
    ext,) (Exp -> (t, Exp)) -> MonoM Exp -> MonoM (t, Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e
    transformAppExp (Loop [VName]
    sparams Pat ParamType
    pat LoopInitBase Info VName
    loopinit LoopFormBase Info VName
    form Exp
    body SrcLoc
    loc) AppRes
    res = do
      Exp
    e1' <- Exp -> MonoM Exp
    transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ LoopInitBase Info VName -> Exp
    loopInitExp LoopInitBase Info VName
    loopinit

      let dimArgs :: Set VName
    dimArgs = [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList [VName]
    sparams
      Pat ParamType
    pat' <- Set VName -> MonoM (Pat ParamType) -> MonoM (Pat ParamType)
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    dimArgs (MonoM (Pat ParamType) -> MonoM (Pat ParamType))
    -> MonoM (Pat ParamType) -> MonoM (Pat ParamType)
    forall a b. (a -> b) -> a -> b
    $ Pat ParamType -> MonoM (Pat ParamType)
    forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
    transformPat Pat ParamType
    pat
      ExpReplacements
    params <- Set VName -> MonoM ExpReplacements
    parametrizing Set VName
    dimArgs
      let sparams' :: [VName]
    sparams' = [VName]
    sparams [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> ((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (ReplacedExp, VName) -> VName
    forall a b. (a, b) -> b
    snd ExpReplacements
    params
          mergeArgs :: Set VName
    mergeArgs = Set VName
    dimArgs Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList (Pat ParamType -> [VName]
    forall t. Pat t -> [VName]
    patNames Pat ParamType
    pat)

      (LoopFormBase Info VName
    form', Set VName
    formArgs) <- case LoopFormBase Info VName
    form of
        For IdentBase Info VName StructType
    ident Exp
    e2 -> (,VName -> Set VName
    forall a. a -> Set a
    S.singleton (VName -> Set VName) -> VName -> Set 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
    identName IdentBase Info VName StructType
    ident) (LoopFormBase Info VName -> (LoopFormBase Info VName, Set VName))
    -> (Exp -> LoopFormBase Info VName)
    -> Exp
    -> (LoopFormBase Info VName, Set VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . IdentBase Info VName StructType -> Exp -> LoopFormBase Info VName
    forall (f :: * -> *) vn.
    IdentBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
    For IdentBase Info VName StructType
    ident (Exp -> (LoopFormBase Info VName, Set VName))
    -> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e2
        ForIn PatBase Info VName StructType
    pat2 Exp
    e2 -> do
          PatBase Info VName StructType
    pat2' <- PatBase Info VName StructType
    -> MonoM (PatBase Info VName StructType)
    forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
    transformPat PatBase Info VName StructType
    pat2
          (,[VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList (PatBase Info VName StructType -> [VName]
    forall t. Pat t -> [VName]
    patNames PatBase Info VName StructType
    pat2)) (LoopFormBase Info VName -> (LoopFormBase Info VName, Set VName))
    -> (Exp -> LoopFormBase Info VName)
    -> Exp
    -> (LoopFormBase Info VName, Set VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatBase Info VName StructType -> Exp -> LoopFormBase Info VName
    forall (f :: * -> *) vn.
    PatBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
    ForIn PatBase Info VName StructType
    pat2' (Exp -> (LoopFormBase Info VName, Set VName))
    -> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e2
        While Exp
    e2 ->
          (Exp -> (LoopFormBase Info VName, Set VName))
    -> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
    forall a b. (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((,Set VName
    forall a. Monoid a => a
    mempty) (LoopFormBase Info VName -> (LoopFormBase Info VName, Set VName))
    -> (Exp -> LoopFormBase Info VName)
    -> Exp
    -> (LoopFormBase Info VName, Set VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Exp -> LoopFormBase Info VName
    forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
    While) (MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName))
    -> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
    forall a b. (a -> b) -> a -> b
    $
            ExpReplacements -> MonoM Exp -> MonoM Exp
    forall a. ExpReplacements -> MonoM a -> MonoM a
    withParams ExpReplacements
    params (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
              Set VName -> MonoM Exp -> MonoM Exp
    scoping Set VName
    mergeArgs (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
                Exp -> MonoM Exp
    transformExp Exp
    e2
      let argset :: Set VName
    argset = Set VName
    mergeArgs Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` Set VName
    formArgs

      Exp
    body' <- ExpReplacements -> MonoM Exp -> MonoM Exp
    forall a. ExpReplacements -> MonoM a -> MonoM a
    withParams ExpReplacements
    params (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Set VName -> MonoM Exp -> MonoM Exp
    scoping Set VName
    argset (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Exp -> MonoM Exp
    transformExp Exp
    body
      -- Maybe monomorphisation introduced new arrays to the loop, and
      -- maybe they have AnySize sizes.  This is not allowed.  Invent some
      -- sizes for them.
      ([VName]
    pat_sizes, Pat ParamType
    pat'') <- Pat ParamType -> MonoM ([VName], Pat ParamType)
    forall (m :: * -> *).
    MonadFreshNames m =>
    Pat ParamType -> m ([VName], Pat ParamType)
    sizesForPat Pat ParamType
    pat'
      AppRes
    res' <- AppRes -> MonoM AppRes
    transformAppRes AppRes
    res
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp ([VName]
    -> Pat ParamType
    -> LoopInitBase Info VName
    -> LoopFormBase Info VName
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [VName]
    -> PatBase f vn ParamType
    -> LoopInitBase f vn
    -> LoopFormBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    Loop ([VName]
    sparams' [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName]
    pat_sizes) Pat ParamType
    pat'' (Exp -> LoopInitBase Info VName
    forall (f :: * -> *) vn. ExpBase f vn -> LoopInitBase f vn
    LoopInitExplicit Exp
    e1') LoopFormBase Info VName
    form' Exp
    body' SrcLoc
    loc) (AppRes -> Info AppRes
    forall a. a -> Info a
    Info AppRes
    res')
    transformAppExp (BinOp (QualName VName
    fname, SrcLoc
    _) (Info StructType
    t) (Exp
    e1, Info (Maybe VName)
    d1) (Exp
    e2, Info (Maybe VName)
    d2) SrcLoc
    loc) AppRes
    res = do
      (AppRes StructType
    ret [VName]
    ext) <- AppRes -> MonoM AppRes
    transformAppRes AppRes
    res
      Exp
    fname' <- SrcLoc -> QualName VName -> StructType -> MonoM Exp
    transformFName SrcLoc
    loc QualName VName
    fname (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct StructType
    t)
      Exp
    e1' <- Exp -> MonoM Exp
    transformExp Exp
    e1
      Exp
    e2' <- Exp -> MonoM Exp
    transformExp Exp
    e2
      if StructType -> Bool
    forall dim as. TypeBase dim as -> Bool
    orderZero (Exp -> StructType
    typeOf Exp
    e1') Bool -> Bool -> Bool
    && StructType -> Bool
    forall dim as. TypeBase dim as -> Bool
    orderZero (Exp -> StructType
    typeOf Exp
    e2')
        then Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ StructType -> [VName] -> Exp -> Exp -> Exp -> Exp
    forall {vn}.
    StructType
    -> [VName]
    -> ExpBase Info vn
    -> ExpBase Info vn
    -> ExpBase Info vn
    -> ExpBase Info vn
    applyOp StructType
    ret [VName]
    ext Exp
    fname' Exp
    e1' Exp
    e2'
        else do
          -- We have to flip the arguments to the function, because
          -- operator application is left-to-right, while function
          -- application is outside-in.  This matters when the arguments
          -- produce existential sizes.  There are later places in the
          -- compiler where we transform BinOp to Apply, but anything that
          -- involves existential sizes will necessarily go through here.
          (Exp
    x_param_e, PatBase Info VName StructType
    x_param) <- Exp -> MonoM (Exp, PatBase Info VName StructType)
    forall {m :: * -> *}.
    MonadFreshNames m =>
    Exp -> m (Exp, PatBase Info VName StructType)
    makeVarParam Exp
    e1'
          (Exp
    y_param_e, PatBase Info VName StructType
    y_param) <- Exp -> MonoM (Exp, PatBase Info VName StructType)
    forall {m :: * -> *}.
    MonadFreshNames m =>
    Exp -> m (Exp, PatBase Info VName StructType)
    makeVarParam Exp
    e2'
          -- XXX: the type annotations here are wrong, but hopefully it
          -- doesn't matter as there will be an outer AppExp to handle
          -- them.
          Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
            AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp
              ( [SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    LetPat
                  []
                  PatBase Info VName StructType
    x_param
                  Exp
    e1'
                  ( AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp
                      ([SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    LetPat [] PatBase Info VName StructType
    y_param Exp
    e2' (StructType -> [VName] -> Exp -> Exp -> Exp -> Exp
    forall {vn}.
    StructType
    -> [VName]
    -> ExpBase Info vn
    -> ExpBase Info vn
    -> ExpBase Info vn
    -> ExpBase Info vn
    applyOp StructType
    ret [VName]
    ext Exp
    fname' Exp
    x_param_e Exp
    y_param_e) SrcLoc
    loc)
                      (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 StructType
    ret [VName]
    forall a. Monoid a => a
    mempty)
                  )
                  SrcLoc
    forall a. Monoid a => a
    mempty
              )
              (AppRes -> Info AppRes
    forall a. a -> Info a
    Info (StructType -> [VName] -> AppRes
    AppRes StructType
    ret [VName]
    forall a. Monoid a => a
    mempty))
      where
        applyOp :: StructType
    -> [VName]
    -> ExpBase Info vn
    -> ExpBase Info vn
    -> ExpBase Info vn
    -> ExpBase Info vn
    applyOp StructType
    ret [VName]
    ext ExpBase Info vn
    fname' ExpBase Info vn
    x ExpBase Info vn
    y =
          ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply
            (ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply ExpBase Info vn
    fname' [(Info (Maybe VName) -> Maybe VName
    forall a. Info a -> a
    unInfo Info (Maybe VName)
    d1, ExpBase Info vn
    x)] (StructType -> [VName] -> AppRes
    AppRes StructType
    ret [VName]
    forall a. Monoid a => a
    mempty))
            [(Info (Maybe VName) -> Maybe VName
    forall a. Info a -> a
    unInfo Info (Maybe VName)
    d2, ExpBase Info vn
    y)]
            (StructType -> [VName] -> AppRes
    AppRes StructType
    ret [VName]
    ext)

        makeVarParam :: Exp -> m (Exp, PatBase Info VName StructType)
    makeVarParam Exp
    arg = do
          let argtype :: StructType
    argtype = Exp -> StructType
    typeOf Exp
    arg
          VName
    x <- [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newNameFromString [Char]
    "binop_p"
          (Exp, PatBase Info VName StructType)
    -> m (Exp, PatBase Info VName StructType)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
            ( QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    x) (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    argtype) SrcLoc
    forall a. Monoid a => a
    mempty,
              VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    x (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    argtype) SrcLoc
    forall a. Monoid a => a
    mempty
            )
    transformAppExp LetWith {} AppRes
    _ =
      [Char] -> MonoM Exp
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "transformAppExp: LetWith is not supposed to occur"
    transformAppExp (Index Exp
    e0 SliceBase Info VName
    idxs SrcLoc
    loc) AppRes
    res =
      AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp
        (AppExpBase Info VName -> Info AppRes -> Exp)
    -> MonoM (AppExpBase Info VName) -> MonoM (Info AppRes -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> SliceBase Info VName -> SrcLoc -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
    Index (Exp -> SliceBase Info VName -> SrcLoc -> AppExpBase Info VName)
    -> MonoM Exp
    -> MonoM (SliceBase Info VName -> SrcLoc -> AppExpBase Info VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e0 MonoM (SliceBase Info VName -> SrcLoc -> AppExpBase Info VName)
    -> MonoM (SliceBase Info VName)
    -> MonoM (SrcLoc -> AppExpBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
    -> SliceBase Info VName -> MonoM (SliceBase Info 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 DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
    transformDimIndex SliceBase Info VName
    idxs MonoM (SrcLoc -> AppExpBase Info VName)
    -> MonoM SrcLoc -> MonoM (AppExpBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc)
        MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (AppRes -> Info AppRes
    forall a. a -> Info a
    Info (AppRes -> Info AppRes) -> MonoM AppRes -> MonoM (Info AppRes)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> AppRes -> MonoM AppRes
    transformAppRes AppRes
    res)
    transformAppExp (Match Exp
    e NonEmpty (CaseBase Info VName)
    cs SrcLoc
    loc) AppRes
    res = do
      Set VName
    implicitDims <- Set VName -> MonoM (Set VName)
    askIntros (Set VName -> MonoM (Set VName)) -> Set VName -> MonoM (Set VName)
    forall a b. (a -> b) -> a -> b
    $ FV -> Set VName
    fvVars (FV -> Set VName) -> FV -> Set VName
    forall a b. (a -> b) -> a -> b
    $ StructType -> FV
    forall u. TypeBase Exp u -> FV
    freeInType (StructType -> FV) -> StructType -> FV
    forall a b. (a -> b) -> a -> b
    $ Exp -> StructType
    typeOf Exp
    e
      Exp
    e' <- Exp -> MonoM Exp
    transformExp Exp
    e
      NonEmpty (CaseBase Info VName)
    cs' <- (CaseBase Info VName -> MonoM (CaseBase Info VName))
    -> NonEmpty (CaseBase Info VName)
    -> MonoM (NonEmpty (CaseBase Info 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) -> NonEmpty a -> m (NonEmpty b)
    mapM (Set VName -> CaseBase Info VName -> MonoM (CaseBase Info VName)
    transformCase Set VName
    implicitDims) NonEmpty (CaseBase Info VName)
    cs
      AppRes
    res' <- AppRes -> MonoM AppRes
    transformAppRes AppRes
    res
      if Set VName -> Bool
    forall a. Set a -> Bool
    S.null Set VName
    implicitDims
        then Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp (Exp
    -> NonEmpty (CaseBase Info VName)
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
    Match Exp
    e' NonEmpty (CaseBase Info VName)
    cs' SrcLoc
    loc) (AppRes -> Info AppRes
    forall a. a -> Info a
    Info AppRes
    res')
        else do
          VName
    tmpVar <- [Char] -> MonoM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newNameFromString [Char]
    "matched_variable"
          Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
            AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp
              ( [SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    LetPat
                  ((VName -> SizeBinder VName) -> [VName] -> [SizeBinder VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> SrcLoc -> SizeBinder VName
    forall vn. vn -> SrcLoc -> SizeBinder vn
    `SizeBinder` SrcLoc
    forall a. Monoid a => a
    mempty) ([VName] -> [SizeBinder VName]) -> [VName] -> [SizeBinder VName]
    forall a b. (a -> b) -> a -> b
    $ Set VName -> [VName]
    forall a. Set a -> [a]
    S.toList Set VName
    implicitDims)
                  (VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    tmpVar (StructType -> Info StructType
    forall a. a -> Info a
    Info (StructType -> Info StructType) -> StructType -> Info StructType
    forall a b. (a -> b) -> a -> b
    $ Exp -> StructType
    typeOf Exp
    e') SrcLoc
    forall a. Monoid a => a
    mempty)
                  Exp
    e'
                  ( AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp
                      (Exp
    -> NonEmpty (CaseBase Info VName)
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
    Match (QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    tmpVar) (StructType -> Info StructType
    forall a. a -> Info a
    Info (StructType -> Info StructType) -> StructType -> Info StructType
    forall a b. (a -> b) -> a -> b
    $ Exp -> StructType
    typeOf Exp
    e') SrcLoc
    forall a. Monoid a => a
    mempty) NonEmpty (CaseBase Info VName)
    cs' SrcLoc
    loc)
                      (AppRes -> Info AppRes
    forall a. a -> Info a
    Info AppRes
    res)
                  )
                  SrcLoc
    forall a. Monoid a => a
    mempty
              )
              (AppRes -> Info AppRes
    forall a. a -> Info a
    Info AppRes
    res')

    -- Monomorphization of expressions.
    transformExp :: Exp -> MonoM Exp
    transformExp :: Exp -> MonoM Exp
    transformExp e :: Exp
    e@Literal {} = Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e
    transformExp e :: Exp
    e@IntLit {} = Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e
    transformExp e :: Exp
    e@FloatLit {} = Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e
    transformExp e :: Exp
    e@StringLit {} = Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e
    transformExp (Parens Exp
    e SrcLoc
    loc) =
      Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
    Parens (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (QualParens (QualName VName, SrcLoc)
    qn Exp
    e SrcLoc
    loc) =
      (QualName VName, SrcLoc) -> Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    (QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
    QualParens (QualName VName, SrcLoc)
    qn (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (TupLit [Exp]
    es SrcLoc
    loc) =
      [Exp] -> SrcLoc -> Exp
    forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
    TupLit ([Exp] -> SrcLoc -> Exp) -> MonoM [Exp] -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [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 -> MonoM Exp
    transformExp [Exp]
    es MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (RecordLit [FieldBase Info VName]
    fs SrcLoc
    loc) =
      [FieldBase Info VName] -> SrcLoc -> Exp
    forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
    RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
    -> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (FieldBase Info VName -> MonoM (FieldBase Info VName))
    -> [FieldBase Info VName] -> MonoM [FieldBase Info 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 FieldBase Info VName -> MonoM (FieldBase Info VName)
    transformField [FieldBase Info VName]
    fs MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
      where
        transformField :: FieldBase Info VName -> MonoM (FieldBase Info VName)
    transformField (RecordFieldExplicit L Name
    name Exp
    e SrcLoc
    loc') =
          L Name -> Exp -> SrcLoc -> FieldBase Info VName
    forall (f :: * -> *) vn.
    L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
    RecordFieldExplicit L Name
    name (Exp -> SrcLoc -> FieldBase Info VName)
    -> MonoM Exp -> MonoM (SrcLoc -> FieldBase Info VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (SrcLoc -> FieldBase Info VName)
    -> MonoM SrcLoc -> MonoM (FieldBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc'
        transformField (RecordFieldImplicit (L Loc
    vloc VName
    v) Info StructType
    t SrcLoc
    _) = do
          Info StructType
    t' <- (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    t
          FieldBase Info VName -> MonoM (FieldBase Info VName)
    transformField (FieldBase Info VName -> MonoM (FieldBase Info VName))
    -> FieldBase Info VName -> MonoM (FieldBase Info VName)
    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
    RecordFieldExplicit
              (Loc -> Name -> L Name
    forall a. Loc -> a -> L a
    L Loc
    vloc (VName -> Name
    baseName VName
    v))
              (QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    v) Info StructType
    t' SrcLoc
    loc)
              SrcLoc
    loc
    transformExp (ArrayVal [PrimValue]
    vs PrimType
    t SrcLoc
    loc) =
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ [PrimValue] -> PrimType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    [PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
    ArrayVal [PrimValue]
    vs PrimType
    t SrcLoc
    loc
    transformExp (ArrayLit [Exp]
    es Info StructType
    t SrcLoc
    loc) =
      [Exp] -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
    ArrayLit ([Exp] -> Info StructType -> SrcLoc -> Exp)
    -> MonoM [Exp] -> MonoM (Info StructType -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [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 -> MonoM Exp
    transformExp [Exp]
    es MonoM (Info StructType -> SrcLoc -> Exp)
    -> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (AppExp AppExpBase Info VName
    e Info AppRes
    res) =
      AppExpBase Info VName -> AppRes -> MonoM Exp
    transformAppExp AppExpBase Info VName
    e (Info AppRes -> AppRes
    forall a. Info a -> a
    unInfo Info AppRes
    res)
    transformExp (Var QualName VName
    fname (Info StructType
    t) SrcLoc
    loc) =
      SrcLoc -> QualName VName -> StructType -> MonoM Exp
    transformFName SrcLoc
    loc QualName VName
    fname (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct StructType
    t)
    transformExp (Hole Info StructType
    t SrcLoc
    loc) =
      Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
    Hole (Info StructType -> SrcLoc -> Exp)
    -> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Ascript Exp
    e TypeExp Exp VName
    tp SrcLoc
    loc) =
      Exp -> TypeExp Exp VName -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn -> TypeExp (ExpBase f vn) vn -> SrcLoc -> ExpBase f vn
    Ascript (Exp -> TypeExp Exp VName -> SrcLoc -> Exp)
    -> MonoM Exp -> MonoM (TypeExp Exp VName -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (TypeExp Exp VName -> SrcLoc -> Exp)
    -> MonoM (TypeExp Exp VName) -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> TypeExp Exp VName -> MonoM (TypeExp Exp VName)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure TypeExp Exp VName
    tp MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Coerce Exp
    e TypeExp Exp VName
    te Info StructType
    t SrcLoc
    loc) =
      Exp -> TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> TypeExp (ExpBase f vn) vn
    -> f StructType
    -> SrcLoc
    -> ExpBase f vn
    Coerce (Exp -> TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp)
    -> MonoM Exp
    -> MonoM (TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp)
    -> MonoM (TypeExp Exp VName)
    -> MonoM (Info StructType -> SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> TypeExp Exp VName -> MonoM (TypeExp Exp VName)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure TypeExp Exp VName
    te MonoM (Info StructType -> SrcLoc -> Exp)
    -> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Negate Exp
    e SrcLoc
    loc) =
      Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
    Negate (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Not Exp
    e SrcLoc
    loc) =
      Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
    Not (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Lambda {}) =
      [Char] -> MonoM Exp
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "transformExp: Lambda is not supposed to occur"
    transformExp (OpSection QualName VName
    qn Info StructType
    t SrcLoc
    loc) =
      Exp -> MonoM Exp
    transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var QualName VName
    qn Info StructType
    t SrcLoc
    loc
    transformExp (OpSectionLeft QualName VName
    fname (Info StructType
    t) Exp
    e (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
    arg (Info RetTypeBase Exp Uniqueness
    rettype, Info [VName]
    retext) SrcLoc
    loc) = do
      let (Info (PName
    xp, ParamType
    xtype, Maybe VName
    xargext), Info (PName
    yp, ParamType
    ytype)) = (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
    arg
      Exp
    e' <- Exp -> MonoM Exp
    transformExp Exp
    e
      QualName VName
    -> Maybe Exp
    -> Maybe Exp
    -> StructType
    -> (PName, ParamType, Maybe VName)
    -> (PName, ParamType, Maybe VName)
    -> (RetTypeBase Exp Uniqueness, [VName])
    -> SrcLoc
    -> MonoM Exp
    desugarBinOpSection
        QualName VName
    fname
        (Exp -> Maybe Exp
    forall a. a -> Maybe a
    Just Exp
    e')
        Maybe Exp
    forall a. Maybe a
    Nothing
        StructType
    t
        (PName
    xp, ParamType
    xtype, Maybe VName
    xargext)
        (PName
    yp, ParamType
    ytype, Maybe VName
    forall a. Maybe a
    Nothing)
        (RetTypeBase Exp Uniqueness
    rettype, [VName]
    retext)
        SrcLoc
    loc
    transformExp (OpSectionRight QualName VName
    fname (Info StructType
    t) Exp
    e (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
    arg (Info RetTypeBase Exp Uniqueness
    rettype) SrcLoc
    loc) = do
      let (Info (PName
    xp, ParamType
    xtype), Info (PName
    yp, ParamType
    ytype, Maybe VName
    yargext)) = (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
    arg
      Exp
    e' <- Exp -> MonoM Exp
    transformExp Exp
    e
      QualName VName
    -> Maybe Exp
    -> Maybe Exp
    -> StructType
    -> (PName, ParamType, Maybe VName)
    -> (PName, ParamType, Maybe VName)
    -> (RetTypeBase Exp Uniqueness, [VName])
    -> SrcLoc
    -> MonoM Exp
    desugarBinOpSection
        QualName VName
    fname
        Maybe Exp
    forall a. Maybe a
    Nothing
        (Exp -> Maybe Exp
    forall a. a -> Maybe a
    Just Exp
    e')
        StructType
    t
        (PName
    xp, ParamType
    xtype, Maybe VName
    forall a. Maybe a
    Nothing)
        (PName
    yp, ParamType
    ytype, Maybe VName
    yargext)
        (RetTypeBase Exp Uniqueness
    rettype, [])
        SrcLoc
    loc
    transformExp (ProjectSection [Name]
    fields (Info StructType
    t) SrcLoc
    loc) = do
      StructType
    t' <- StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    t
      [Name] -> StructType -> SrcLoc -> MonoM Exp
    desugarProjectSection [Name]
    fields StructType
    t' SrcLoc
    loc
    transformExp (IndexSection SliceBase Info VName
    idxs (Info StructType
    t) SrcLoc
    loc) = do
      SliceBase Info VName
    idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
    -> SliceBase Info VName -> MonoM (SliceBase Info 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 DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
    transformDimIndex SliceBase Info VName
    idxs
      SliceBase Info VName -> StructType -> SrcLoc -> MonoM Exp
    desugarIndexSection SliceBase Info VName
    idxs' StructType
    t SrcLoc
    loc
    transformExp (Project Name
    n Exp
    e Info StructType
    tp SrcLoc
    loc) = do
      Info StructType
    tp' <- (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    tp
      Exp
    e' <- Exp -> MonoM Exp
    transformExp Exp
    e
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Name -> Exp -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
    Project Name
    n Exp
    e' Info StructType
    tp' SrcLoc
    loc
    transformExp (Update Exp
    e1 SliceBase Info VName
    idxs Exp
    e2 SrcLoc
    loc) =
      Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
    Update
        (Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp)
    -> MonoM Exp
    -> MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e1
        MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
    -> MonoM (SliceBase Info VName) -> MonoM (Exp -> SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
    -> SliceBase Info VName -> MonoM (SliceBase Info 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 DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
    transformDimIndex SliceBase Info VName
    idxs
        MonoM (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Exp -> MonoM Exp
    transformExp Exp
    e2
        MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (RecordUpdate Exp
    e1 [Name]
    fs Exp
    e2 Info StructType
    t SrcLoc
    loc) =
      Exp -> [Name] -> Exp -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn
    -> [Name] -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
    RecordUpdate
        (Exp -> [Name] -> Exp -> Info StructType -> SrcLoc -> Exp)
    -> MonoM Exp
    -> MonoM ([Name] -> Exp -> Info StructType -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e1
        MonoM ([Name] -> Exp -> Info StructType -> SrcLoc -> Exp)
    -> MonoM [Name] -> MonoM (Exp -> Info StructType -> SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [Name] -> MonoM [Name]
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Name]
    fs
        MonoM (Exp -> Info StructType -> SrcLoc -> Exp)
    -> MonoM Exp -> MonoM (Info StructType -> SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Exp -> MonoM Exp
    transformExp Exp
    e2
        MonoM (Info StructType -> SrcLoc -> Exp)
    -> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    t
        MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Assert Exp
    e1 Exp
    e2 Info Text
    desc SrcLoc
    loc) =
      Exp -> Exp -> Info Text -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
    Assert (Exp -> Exp -> Info Text -> SrcLoc -> Exp)
    -> MonoM Exp -> MonoM (Exp -> Info Text -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e1 MonoM (Exp -> Info Text -> SrcLoc -> Exp)
    -> MonoM Exp -> MonoM (Info Text -> SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Exp -> MonoM Exp
    transformExp Exp
    e2 MonoM (Info Text -> SrcLoc -> Exp)
    -> MonoM (Info Text) -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Info Text -> MonoM (Info Text)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Info Text
    desc MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Constr Name
    name [Exp]
    all_es Info StructType
    t SrcLoc
    loc) =
      Name -> [Exp] -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    Name -> [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
    Constr Name
    name ([Exp] -> Info StructType -> SrcLoc -> Exp)
    -> MonoM [Exp] -> MonoM (Info StructType -> SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [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 -> MonoM Exp
    transformExp [Exp]
    all_es MonoM (Info StructType -> SrcLoc -> Exp)
    -> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (StructType -> MonoM StructType)
    -> Info StructType -> MonoM (Info StructType)
    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) -> Info a -> f (Info b)
    traverse StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType Info StructType
    t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc
    transformExp (Attr AttrInfo VName
    info Exp
    e SrcLoc
    loc) =
      AttrInfo VName -> Exp -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
    Attr AttrInfo VName
    info (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc

    transformCase :: S.Set VName -> Case -> MonoM Case
    transformCase :: Set VName -> CaseBase Info VName -> MonoM (CaseBase Info VName)
    transformCase Set VName
    implicitDims (CasePat PatBase Info VName StructType
    p Exp
    e SrcLoc
    loc) = do
      PatBase Info VName StructType
    p' <- PatBase Info VName StructType
    -> MonoM (PatBase Info VName StructType)
    forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
    transformPat PatBase Info VName StructType
    p
      PatBase Info VName StructType
    -> Exp -> SrcLoc -> CaseBase Info VName
    forall (f :: * -> *) vn.
    PatBase f vn StructType -> ExpBase f vn -> SrcLoc -> CaseBase f vn
    CasePat PatBase Info VName StructType
    p' (Exp -> SrcLoc -> CaseBase Info VName)
    -> MonoM Exp -> MonoM (SrcLoc -> CaseBase Info VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Set VName -> MonoM Exp -> MonoM Exp
    scoping ([VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList (PatBase Info VName StructType -> [VName]
    forall t. Pat t -> [VName]
    patNames PatBase Info VName StructType
    p) Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` Set VName
    implicitDims) (Exp -> MonoM Exp
    transformExp Exp
    e) MonoM (SrcLoc -> CaseBase Info VName)
    -> MonoM SrcLoc -> MonoM (CaseBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SrcLoc -> MonoM SrcLoc
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SrcLoc
    loc

    transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
    transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
    transformDimIndex (DimFix Exp
    e) = Exp -> DimIndexBase Info VName
    forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
    DimFix (Exp -> DimIndexBase Info VName)
    -> MonoM Exp -> MonoM (DimIndexBase Info VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> MonoM Exp
    transformExp Exp
    e
    transformDimIndex (DimSlice Maybe Exp
    me1 Maybe Exp
    me2 Maybe Exp
    me3) =
      Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndexBase Info VName
    forall (f :: * -> *) vn.
    Maybe (ExpBase f vn)
    -> Maybe (ExpBase f vn)
    -> Maybe (ExpBase f vn)
    -> DimIndexBase f vn
    DimSlice (Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
    -> MonoM (Maybe Exp)
    -> MonoM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Maybe Exp -> MonoM (Maybe Exp)
    trans Maybe Exp
    me1 MonoM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
    -> MonoM (Maybe Exp)
    -> MonoM (Maybe Exp -> DimIndexBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Maybe Exp -> MonoM (Maybe Exp)
    trans Maybe Exp
    me2 MonoM (Maybe Exp -> DimIndexBase Info VName)
    -> MonoM (Maybe Exp) -> MonoM (DimIndexBase Info VName)
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Maybe Exp -> MonoM (Maybe Exp)
    trans Maybe Exp
    me3
      where
        trans :: Maybe Exp -> MonoM (Maybe Exp)
    trans = (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe 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) -> Maybe a -> m (Maybe b)
    mapM Exp -> MonoM Exp
    transformExp

    -- Transform an operator section into a lambda.
    desugarBinOpSection ::
      QualName VName ->
      Maybe Exp ->
      Maybe Exp ->
      StructType ->
      (PName, ParamType, Maybe VName) ->
      (PName, ParamType, Maybe VName) ->
      (ResRetType, [VName]) ->
      SrcLoc ->
      MonoM Exp
    desugarBinOpSection :: QualName VName
    -> Maybe Exp
    -> Maybe Exp
    -> StructType
    -> (PName, ParamType, Maybe VName)
    -> (PName, ParamType, Maybe VName)
    -> (RetTypeBase Exp Uniqueness, [VName])
    -> SrcLoc
    -> MonoM Exp
    desugarBinOpSection QualName VName
    fname Maybe Exp
    e_left Maybe Exp
    e_right StructType
    t (PName
    xp, ParamType
    xtype, Maybe VName
    xext) (PName
    yp, ParamType
    ytype, Maybe VName
    yext) (RetType [VName]
    dims TypeBase Exp Uniqueness
    rettype, [VName]
    retext) SrcLoc
    loc = do
      StructType
    t' <- StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    t
      Exp
    op <- SrcLoc -> QualName VName -> StructType -> MonoM Exp
    transformFName SrcLoc
    loc QualName VName
    fname (StructType -> MonoM Exp) -> StructType -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct StructType
    t
      (VName
    v1, Exp -> Exp
    wrap_left, Exp
    e1, [Pat ParamType]
    p1) <- Maybe Exp
    -> ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
    forall {m :: * -> *} {u}.
    MonadFreshNames m =>
    Maybe Exp
    -> TypeBase Exp u
    -> m (VName, Exp -> Exp, Exp,
          [PatBase Info VName (TypeBase Exp u)])
    makeVarParam Maybe Exp
    e_left (ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType]))
    -> MonoM ParamType
    -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< ParamType -> MonoM ParamType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType ParamType
    xtype
      (VName
    v2, Exp -> Exp
    wrap_right, Exp
    e2, [Pat ParamType]
    p2) <- Maybe Exp
    -> ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
    forall {m :: * -> *} {u}.
    MonadFreshNames m =>
    Maybe Exp
    -> TypeBase Exp u
    -> m (VName, Exp -> Exp, Exp,
          [PatBase Info VName (TypeBase Exp u)])
    makeVarParam Maybe Exp
    e_right (ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType]))
    -> MonoM ParamType
    -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< ParamType -> MonoM ParamType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType ParamType
    ytype
      let apply_left :: Exp
    apply_left =
            Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply
              Exp
    op
              [(Maybe VName
    xext, Exp
    e1)]
              (StructType -> [VName] -> AppRes
    AppRes (ScalarTypeBase Exp NoUniqueness -> StructType
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
    -> ScalarTypeBase Exp NoUniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $ NoUniqueness
    -> PName
    -> Diet
    -> StructType
    -> RetTypeBase Exp Uniqueness
    -> ScalarTypeBase Exp NoUniqueness
    forall dim u.
    u
    -> PName
    -> Diet
    -> TypeBase dim NoUniqueness
    -> RetTypeBase dim Uniqueness
    -> ScalarTypeBase dim u
    Arrow NoUniqueness
    forall a. Monoid a => a
    mempty PName
    yp (ParamType -> Diet
    forall shape. TypeBase shape Diet -> Diet
    diet ParamType
    ytype) (ParamType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct ParamType
    ytype) ([VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [] (TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness)
    -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall a b. (a -> b) -> a -> b
    $ Uniqueness -> StructType -> TypeBase Exp Uniqueness
    forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
    toRes Uniqueness
    Nonunique StructType
    t')) [])
          onDim :: ExpBase f VName -> ExpBase f VName
    onDim (Var QualName VName
    d f StructType
    typ SrcLoc
    _)
            | Named VName
    p <- PName
    xp, QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    d VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    p = QualName VName -> f StructType -> SrcLoc -> ExpBase f VName
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    v1) f StructType
    typ SrcLoc
    loc
            | Named VName
    p <- PName
    yp, QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    d VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    p = QualName VName -> f StructType -> SrcLoc -> ExpBase f VName
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    v2) f StructType
    typ SrcLoc
    loc
          onDim ExpBase f VName
    d = ExpBase f VName
    d
          rettype' :: TypeBase Exp Uniqueness
    rettype' = (Exp -> Exp) -> TypeBase Exp Uniqueness -> TypeBase Exp Uniqueness
    forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
    forall (p :: * -> * -> *) a b c.
    Bifunctor p =>
    (a -> b) -> p a c -> p b c
    first Exp -> Exp
    forall {f :: * -> *}. ExpBase f VName -> ExpBase f VName
    onDim TypeBase Exp Uniqueness
    rettype
      Exp
    body <-
        Set VName -> MonoM Exp -> MonoM Exp
    scoping ([VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList [VName
    v1, VName
    v2]) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
          Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
    forall vn.
    ExpBase Info vn
    -> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
    mkApply Exp
    apply_left [(Maybe VName
    yext, Exp
    e2)]
            (AppRes -> Exp) -> MonoM AppRes -> MonoM Exp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> AppRes -> MonoM AppRes
    transformAppRes (StructType -> [VName] -> AppRes
    AppRes (TypeBase Exp Uniqueness -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct TypeBase Exp Uniqueness
    rettype') [VName]
    retext)
      RetTypeBase Exp Uniqueness
    rettype'' <- Set VName
    -> RetTypeBase Exp Uniqueness -> MonoM (RetTypeBase Exp Uniqueness)
    forall as.
    Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    transformRetTypeSizes ([VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList [VName
    v1, VName
    v2]) (RetTypeBase Exp Uniqueness -> MonoM (RetTypeBase Exp Uniqueness))
    -> RetTypeBase Exp Uniqueness -> MonoM (RetTypeBase Exp Uniqueness)
    forall a b. (a -> b) -> a -> b
    $ [VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [VName]
    dims TypeBase Exp Uniqueness
    rettype'
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> (Exp -> Exp) -> Exp -> MonoM Exp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Exp -> Exp
    wrap_left (Exp -> Exp) -> (Exp -> Exp) -> Exp -> Exp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Exp -> Exp
    wrap_right (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
        [Pat ParamType]
    -> Exp
    -> Maybe (TypeExp Exp VName)
    -> Info (RetTypeBase Exp Uniqueness)
    -> SrcLoc
    -> Exp
    forall (f :: * -> *) vn.
    [PatBase f vn ParamType]
    -> ExpBase f vn
    -> Maybe (TypeExp (ExpBase f vn) vn)
    -> f (RetTypeBase Exp Uniqueness)
    -> SrcLoc
    -> ExpBase f vn
    Lambda ([Pat ParamType]
    p1 [Pat ParamType] -> [Pat ParamType] -> [Pat ParamType]
    forall a. [a] -> [a] -> [a]
    ++ [Pat ParamType]
    p2) Exp
    body Maybe (TypeExp Exp VName)
    forall a. Maybe a
    Nothing (RetTypeBase Exp Uniqueness -> Info (RetTypeBase Exp Uniqueness)
    forall a. a -> Info a
    Info RetTypeBase Exp Uniqueness
    rettype'') SrcLoc
    loc
      where
        patAndVar :: TypeBase Exp u
    -> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
    patAndVar TypeBase Exp u
    argtype = do
          VName
    x <- [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newNameFromString [Char]
    "x"
          (VName, PatBase Info VName (TypeBase Exp u), Exp)
    -> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
            ( VName
    x,
              VName
    -> Info (TypeBase Exp u)
    -> SrcLoc
    -> PatBase Info VName (TypeBase Exp u)
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    x (TypeBase Exp u -> Info (TypeBase Exp u)
    forall a. a -> Info a
    Info TypeBase Exp u
    argtype) SrcLoc
    forall a. Monoid a => a
    mempty,
              QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    x) (StructType -> Info StructType
    forall a. a -> Info a
    Info (TypeBase Exp u -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct TypeBase Exp u
    argtype)) SrcLoc
    forall a. Monoid a => a
    mempty
            )

        makeVarParam :: Maybe Exp
    -> TypeBase Exp u
    -> m (VName, Exp -> Exp, Exp,
          [PatBase Info VName (TypeBase Exp u)])
    makeVarParam (Just Exp
    e) TypeBase Exp u
    argtype = do
          (VName
    v, PatBase Info VName (TypeBase Exp u)
    pat, Exp
    var_e) <- TypeBase Exp u
    -> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
    forall {m :: * -> *} {u}.
    MonadFreshNames m =>
    TypeBase Exp u
    -> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
    patAndVar TypeBase Exp u
    argtype
          let wrap :: Exp -> Exp
    wrap Exp
    body =
                AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp ([SizeBinder VName]
    -> PatBase Info VName StructType
    -> Exp
    -> Exp
    -> SrcLoc
    -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    [SizeBinder vn]
    -> PatBase f vn StructType
    -> ExpBase f vn
    -> ExpBase f vn
    -> SrcLoc
    -> AppExpBase f vn
    LetPat [] ((TypeBase Exp u -> StructType)
    -> PatBase Info VName (TypeBase Exp u)
    -> PatBase Info VName StructType
    forall a b.
    (a -> b) -> PatBase Info VName a -> PatBase Info VName b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap TypeBase Exp u -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct PatBase Info VName (TypeBase Exp u)
    pat) Exp
    e Exp
    body 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 (Exp -> StructType
    typeOf Exp
    body) [VName]
    forall a. Monoid a => a
    mempty)
          (VName, Exp -> Exp, Exp, [PatBase Info VName (TypeBase Exp u)])
    -> m (VName, Exp -> Exp, Exp,
          [PatBase Info VName (TypeBase Exp u)])
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName
    v, Exp -> Exp
    wrap, Exp
    var_e, [])
        makeVarParam Maybe Exp
    Nothing TypeBase Exp u
    argtype = do
          (VName
    v, PatBase Info VName (TypeBase Exp u)
    pat, Exp
    var_e) <- TypeBase Exp u
    -> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
    forall {m :: * -> *} {u}.
    MonadFreshNames m =>
    TypeBase Exp u
    -> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
    patAndVar TypeBase Exp u
    argtype
          (VName, Exp -> Exp, Exp, [PatBase Info VName (TypeBase Exp u)])
    -> m (VName, Exp -> Exp, Exp,
          [PatBase Info VName (TypeBase Exp u)])
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName
    v, Exp -> Exp
    forall a. a -> a
    id, Exp
    var_e, [PatBase Info VName (TypeBase Exp u)
    pat])

    desugarProjectSection :: [Name] -> StructType -> SrcLoc -> MonoM Exp
    desugarProjectSection :: [Name] -> StructType -> SrcLoc -> MonoM Exp
    desugarProjectSection [Name]
    fields (Scalar (Arrow NoUniqueness
    _ PName
    _ Diet
    _ StructType
    t1 (RetType [VName]
    dims TypeBase Exp Uniqueness
    t2))) SrcLoc
    loc = do
      VName
    p <- [Char] -> MonoM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "project_p"
      let body :: Exp
    body = (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Exp -> Name -> Exp
    project (QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    p) (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    t1) SrcLoc
    forall a. Monoid a => a
    mempty) [Name]
    fields
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
        [Pat ParamType]
    -> Exp
    -> Maybe (TypeExp Exp VName)
    -> Info (RetTypeBase Exp Uniqueness)
    -> SrcLoc
    -> Exp
    forall (f :: * -> *) vn.
    [PatBase f vn ParamType]
    -> ExpBase f vn
    -> Maybe (TypeExp (ExpBase f vn) vn)
    -> f (RetTypeBase Exp Uniqueness)
    -> SrcLoc
    -> ExpBase f vn
    Lambda
          [VName -> Info ParamType -> SrcLoc -> Pat ParamType
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    p (ParamType -> Info ParamType
    forall a. a -> Info a
    Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
    forall a b. (a -> b) -> a -> b
    $ Diet -> StructType -> ParamType
    forall u. Diet -> TypeBase Exp u -> ParamType
    toParam Diet
    Observe StructType
    t1) SrcLoc
    forall a. Monoid a => a
    mempty]
          Exp
    body
          Maybe (TypeExp Exp VName)
    forall a. Maybe a
    Nothing
          (RetTypeBase Exp Uniqueness -> Info (RetTypeBase Exp Uniqueness)
    forall a. a -> Info a
    Info ([VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [VName]
    dims TypeBase Exp Uniqueness
    t2))
          SrcLoc
    loc
      where
        project :: Exp -> Name -> Exp
    project Exp
    e Name
    field =
          case Exp -> StructType
    typeOf Exp
    e of
            Scalar (Record Map Name StructType
    fs)
              | Just StructType
    t <- Name -> Map Name StructType -> Maybe StructType
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Name
    field Map Name StructType
    fs ->
                  Name -> Exp -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
    Project Name
    field Exp
    e (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    t) SrcLoc
    forall a. Monoid a => a
    mempty
            StructType
    t ->
              [Char] -> Exp
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> Exp) -> [Char] -> Exp
    forall a b. (a -> b) -> a -> b
    $
                [Char]
    "desugarOpSection: type "
                  [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ StructType -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString StructType
    t
                  [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " does not have field "
                  [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ Name -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString Name
    field
    desugarProjectSection [Name]
    _ StructType
    t SrcLoc
    _ = [Char] -> MonoM Exp
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> MonoM Exp) -> [Char] -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "desugarOpSection: not a function type: " [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ StructType -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString StructType
    t

    desugarIndexSection :: [DimIndex] -> StructType -> SrcLoc -> MonoM Exp
    desugarIndexSection :: SliceBase Info VName -> StructType -> SrcLoc -> MonoM Exp
    desugarIndexSection SliceBase Info VName
    idxs (Scalar (Arrow NoUniqueness
    _ PName
    _ Diet
    _ StructType
    t1 (RetType [VName]
    dims TypeBase Exp Uniqueness
    t2))) SrcLoc
    loc = do
      VName
    p <- [Char] -> MonoM VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "index_i"
      StructType
    t1' <- StructType -> MonoM StructType
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType StructType
    t1
      TypeBase Exp Uniqueness
    t2' <- TypeBase Exp Uniqueness -> MonoM (TypeBase Exp Uniqueness)
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType TypeBase Exp Uniqueness
    t2
      let body :: Exp
    body = AppExpBase Info VName -> Info AppRes -> Exp
    forall (f :: * -> *) vn.
    AppExpBase f vn -> f AppRes -> ExpBase f vn
    AppExp (Exp -> SliceBase Info VName -> SrcLoc -> AppExpBase Info VName
    forall (f :: * -> *) vn.
    ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
    Index (QualName VName -> Info StructType -> SrcLoc -> Exp
    forall (f :: * -> *) vn.
    QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
    Var (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    p) (StructType -> Info StructType
    forall a. a -> Info a
    Info StructType
    t1') SrcLoc
    loc) SliceBase Info VName
    idxs SrcLoc
    loc) (AppRes -> Info AppRes
    forall a. a -> Info a
    Info (StructType -> [VName] -> AppRes
    AppRes (TypeBase Exp Uniqueness -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct TypeBase Exp Uniqueness
    t2') []))
      Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $
        [Pat ParamType]
    -> Exp
    -> Maybe (TypeExp Exp VName)
    -> Info (RetTypeBase Exp Uniqueness)
    -> SrcLoc
    -> Exp
    forall (f :: * -> *) vn.
    [PatBase f vn ParamType]
    -> ExpBase f vn
    -> Maybe (TypeExp (ExpBase f vn) vn)
    -> f (RetTypeBase Exp Uniqueness)
    -> SrcLoc
    -> ExpBase f vn
    Lambda
          [VName -> Info ParamType -> SrcLoc -> Pat ParamType
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    p (ParamType -> Info ParamType
    forall a. a -> Info a
    Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
    forall a b. (a -> b) -> a -> b
    $ Diet -> StructType -> ParamType
    forall u. Diet -> TypeBase Exp u -> ParamType
    toParam Diet
    Observe StructType
    t1') SrcLoc
    forall a. Monoid a => a
    mempty]
          Exp
    body
          Maybe (TypeExp Exp VName)
    forall a. Maybe a
    Nothing
          (RetTypeBase Exp Uniqueness -> Info (RetTypeBase Exp Uniqueness)
    forall a. a -> Info a
    Info ([VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [VName]
    dims TypeBase Exp Uniqueness
    t2'))
          SrcLoc
    loc
    desugarIndexSection SliceBase Info VName
    _ StructType
    t SrcLoc
    _ = [Char] -> MonoM Exp
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> MonoM Exp) -> [Char] -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "desugarIndexSection: not a function type: " [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ StructType -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString StructType
    t

    transformPat :: Pat (TypeBase Size u) -> MonoM (Pat (TypeBase Size u))
    transformPat :: forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
    transformPat = (TypeBase Exp u -> MonoM (TypeBase Exp u))
    -> PatBase Info VName (TypeBase Exp u)
    -> MonoM (PatBase Info VName (TypeBase Exp u))
    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) -> PatBase Info VName a -> f (PatBase Info VName b)
    traverse TypeBase Exp u -> MonoM (TypeBase Exp u)
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType

    type DimInst = M.Map VName Size

    dimMapping ::
      (Monoid a) =>
      TypeBase Size a ->
      TypeBase Size a ->
      ExpReplacements ->
      ExpReplacements ->
      DimInst
    dimMapping :: forall a.
    Monoid a =>
    TypeBase Exp a
    -> TypeBase Exp a -> ExpReplacements -> ExpReplacements -> DimInst
    dimMapping TypeBase Exp a
    t1 TypeBase Exp a
    t2 ExpReplacements
    r1 ExpReplacements
    r2 = State DimInst (TypeBase Exp a) -> DimInst -> DimInst
    forall s a. State s a -> s -> s
    execState (([VName] -> Exp -> Exp -> StateT DimInst Identity Exp)
    -> TypeBase Exp a
    -> TypeBase Exp a
    -> State DimInst (TypeBase Exp a)
    forall as (m :: * -> *) d1 d2.
    (Monoid as, Monad m) =>
    ([VName] -> d1 -> d2 -> m d1)
    -> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
    matchDims [VName] -> Exp -> Exp -> StateT DimInst Identity Exp
    forall {m :: * -> *} {t :: * -> *}.
    (Foldable t, MonadState DimInst m, Monoid (t VName)) =>
    t VName -> Exp -> Exp -> m Exp
    onDims TypeBase Exp a
    t1 TypeBase Exp a
    t2) DimInst
    forall a. Monoid a => a
    mempty
      where
        revMap :: [(b, a)] -> [(a, b)]
    revMap = ((b, a) -> (a, b)) -> [(b, a)] -> [(a, b)]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(b
    k, a
    v) -> (a
    v, b
    k))
        named1 :: [(VName, ReplacedExp)]
    named1 = ExpReplacements -> [(VName, ReplacedExp)]
    forall {b} {a}. [(b, a)] -> [(a, b)]
    revMap ExpReplacements
    r1
        named2 :: [(VName, ReplacedExp)]
    named2 = ExpReplacements -> [(VName, ReplacedExp)]
    forall {b} {a}. [(b, a)] -> [(a, b)]
    revMap ExpReplacements
    r2

        onDims :: t VName -> Exp -> Exp -> m Exp
    onDims t VName
    bound Exp
    e1 Exp
    e2 = do
          t VName -> Exp -> Exp -> m ()
    forall {m :: * -> *} {t :: * -> *}.
    (Foldable t, MonadState DimInst m, Monoid (t VName)) =>
    t VName -> Exp -> Exp -> m ()
    onExps t VName
    bound Exp
    e1 Exp
    e2
          Exp -> m Exp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e1

        -- XXX: It is intentional that we throw away the 'bound'
        -- information after looking up ExpReplacements, as there are
        -- cases (particularly including the function types that occur
        -- after lambda lifting) where some troublesome shadowing occurs,
        -- and the names are not actually locally bound. We cannot just
        -- ignore the bound information entirely, and expect that any
        -- instantiation uses only names in scope at the outer level, due
        -- to truly exotic cases like entry-lifted.fut, where the real
        -- sizes were not visible to the type checker. Arguably that is
        -- the thing that should be fixed, but it requires fiddling with
        -- the defunctorisation of size-lifted types.

        onExps :: t VName -> Exp -> Exp -> m ()
    onExps t VName
    bound (Var QualName VName
    v Info StructType
    _ SrcLoc
    _) Exp
    e = do
          Bool -> m () -> m ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ((VName -> Bool) -> Set VName -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> t VName -> Bool
    forall a. Eq a => a -> t a -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` t VName
    bound) (Set VName -> Bool) -> Set VName -> Bool
    forall a b. (a -> b) -> a -> b
    $ Exp -> Set VName
    freeVarsInExp Exp
    e) (m () -> m ()) -> m () -> m ()
    forall a b. (a -> b) -> a -> b
    $
            (DimInst -> DimInst) -> m ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify (VName -> Exp -> DimInst -> DimInst
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    v) Exp
    e)
          case VName -> [(VName, ReplacedExp)] -> Maybe ReplacedExp
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    v) [(VName, ReplacedExp)]
    named1 of
            Just ReplacedExp
    rexp -> t VName -> Exp -> Exp -> m ()
    onExps t VName
    forall a. Monoid a => a
    mempty (ReplacedExp -> Exp
    unReplaced ReplacedExp
    rexp) Exp
    e
            Maybe ReplacedExp
    Nothing -> () -> m ()
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        onExps t VName
    _bound Exp
    e (Var QualName VName
    v Info StructType
    _ SrcLoc
    _)
          | Just ReplacedExp
    rexp <- VName -> [(VName, ReplacedExp)] -> Maybe ReplacedExp
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    v) [(VName, ReplacedExp)]
    named2 =
              t VName -> Exp -> Exp -> m ()
    onExps t VName
    forall a. Monoid a => a
    mempty Exp
    e (ReplacedExp -> Exp
    unReplaced ReplacedExp
    rexp)
        onExps t VName
    bound Exp
    e1 Exp
    e2
          | Just [(Exp, Exp)]
    es <- Exp -> Exp -> Maybe [(Exp, Exp)]
    similarExps Exp
    e1 Exp
    e2 =
              ((Exp, Exp) -> m ()) -> [(Exp, Exp)] -> m ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ((Exp -> Exp -> m ()) -> (Exp, Exp) -> m ()
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ((Exp -> Exp -> m ()) -> (Exp, Exp) -> m ())
    -> (Exp -> Exp -> m ()) -> (Exp, Exp) -> m ()
    forall a b. (a -> b) -> a -> b
    $ t VName -> Exp -> Exp -> m ()
    onExps t VName
    bound) [(Exp, Exp)]
    es
        onExps t VName
    _ Exp
    _ Exp
    _ = () -> m ()
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    forall a. Monoid a => a
    mempty

        freeVarsInExp :: Exp -> Set VName
    freeVarsInExp = FV -> Set VName
    fvVars (FV -> Set VName) -> (Exp -> FV) -> Exp -> Set VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Exp -> FV
    freeInExp

    inferSizeArgs :: [TypeParam] -> StructType -> ExpReplacements -> StructType -> MonoM [Exp]
    inferSizeArgs :: [TypeParam] -> StructType -> ExpReplacements -> InferSizeArgs
    inferSizeArgs [TypeParam]
    tparams StructType
    bind_t ExpReplacements
    bind_r StructType
    t = do
      ExpReplacements
    r <- (ExpReplacements -> ExpReplacements -> ExpReplacements)
    -> MonoM (ExpReplacements -> ExpReplacements)
    forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
    gets ExpReplacements -> ExpReplacements -> ExpReplacements
    forall a. Semigroup a => a -> a -> a
    (<>) MonoM (ExpReplacements -> ExpReplacements)
    -> MonoM ExpReplacements -> MonoM ExpReplacements
    forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (Env -> ExpReplacements) -> MonoM ExpReplacements
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
    asks Env -> ExpReplacements
    envParametrized
      let dinst :: DimInst
    dinst = StructType
    -> StructType -> ExpReplacements -> ExpReplacements -> DimInst
    forall a.
    Monoid a =>
    TypeBase Exp a
    -> TypeBase Exp a -> ExpReplacements -> ExpReplacements -> DimInst
    dimMapping StructType
    bind_t StructType
    t ExpReplacements
    bind_r ExpReplacements
    r
      (TypeParam -> MonoM Exp) -> [TypeParam] -> MonoM [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 (DimInst -> TypeParam -> MonoM Exp
    forall {k}. Ord k => Map k Exp -> TypeParamBase k -> MonoM Exp
    tparamArg DimInst
    dinst) [TypeParam]
    tparams
      where
        tparamArg :: Map k Exp -> TypeParamBase k -> MonoM Exp
    tparamArg Map k Exp
    dinst TypeParamBase k
    tp =
          case k -> Map k Exp -> Maybe Exp
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup (TypeParamBase k -> k
    forall vn. TypeParamBase vn -> vn
    typeParamName TypeParamBase k
    tp) Map k Exp
    dinst of
            Just Exp
    e
              -- In some cases we infer anySizes for size arguments. This
              -- only occurs when those sizes don't actually matter (knock
              -- on wood...), but we should never actually insert anySize
              -- as a concrete argument.
              | Exp
    e Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Exp
    anySize ->
                  Exp -> MonoM Exp
    replaceExp Exp
    e
            Maybe Exp
    _ ->
              Exp -> MonoM Exp
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Integer -> SrcLoc -> Exp
    sizeFromInteger Integer
    0 SrcLoc
    forall a. Monoid a => a
    mempty

    -- Monomorphising higher-order functions can result in function types
    -- where the same named parameter occurs in multiple spots.  When
    -- monomorphising we don't really need those parameter names anymore,
    -- and the defunctionaliser can be confused if there are duplicates
    -- (it doesn't handle shadowing), so let's just remove all parameter
    -- names here.  This is safe because a MonoType does not contain sizes
    -- anyway.
    noNamedParams :: MonoType -> MonoType
    noNamedParams :: MonoType -> MonoType
    noNamedParams = MonoType -> MonoType
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    f
      where
        f :: TypeBase MonoSize u -> TypeBase MonoSize u
        f :: forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    f (Array u
    u Shape MonoSize
    shape ScalarTypeBase MonoSize NoUniqueness
    t) = u
    -> Shape MonoSize
    -> ScalarTypeBase MonoSize NoUniqueness
    -> TypeBase MonoSize u
    forall dim u.
    u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
    Array u
    u Shape MonoSize
    shape (ScalarTypeBase MonoSize NoUniqueness
    -> ScalarTypeBase MonoSize NoUniqueness
    forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    f' ScalarTypeBase MonoSize NoUniqueness
    t)
        f (Scalar ScalarTypeBase MonoSize u
    t) = ScalarTypeBase MonoSize u -> TypeBase MonoSize u
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase MonoSize u -> TypeBase MonoSize u)
    -> ScalarTypeBase MonoSize u -> TypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    f' ScalarTypeBase MonoSize u
    t
        f' :: ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
        f' :: forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    f' (Record Map Name (TypeBase MonoSize u)
    fs) = Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    Record (Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u)
    -> Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ (TypeBase MonoSize u -> TypeBase MonoSize u)
    -> Map Name (TypeBase MonoSize u) -> Map Name (TypeBase MonoSize u)
    forall a b. (a -> b) -> Map Name a -> Map Name b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap TypeBase MonoSize u -> TypeBase MonoSize u
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    f Map Name (TypeBase MonoSize u)
    fs
        f' (Sum Map Name [TypeBase MonoSize u]
    cs) = Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
    forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
    Sum (Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u)
    -> Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
    forall a b. (a -> b) -> a -> b
    $ ([TypeBase MonoSize u] -> [TypeBase MonoSize u])
    -> Map Name [TypeBase MonoSize u] -> Map Name [TypeBase MonoSize u]
    forall a b. (a -> b) -> Map Name a -> Map Name b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((TypeBase MonoSize u -> TypeBase MonoSize u)
    -> [TypeBase MonoSize u] -> [TypeBase MonoSize u]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase MonoSize u -> TypeBase MonoSize u
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    f) Map Name [TypeBase MonoSize u]
    cs
        f' (Arrow u
    u PName
    _ Diet
    d1 MonoType
    t1 (RetType [VName]
    dims TypeBase MonoSize Uniqueness
    t2)) =
          u
    -> PName
    -> Diet
    -> MonoType
    -> RetTypeBase MonoSize Uniqueness
    -> ScalarTypeBase MonoSize u
    forall dim u.
    u
    -> PName
    -> Diet
    -> TypeBase dim NoUniqueness
    -> RetTypeBase dim Uniqueness
    -> ScalarTypeBase dim u
    Arrow u
    u PName
    Unnamed Diet
    d1 (MonoType -> MonoType
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    f MonoType
    t1) ([VName]
    -> TypeBase MonoSize Uniqueness -> RetTypeBase MonoSize Uniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [VName]
    dims (TypeBase MonoSize Uniqueness -> TypeBase MonoSize Uniqueness
    forall u. TypeBase MonoSize u -> TypeBase MonoSize u
    f TypeBase MonoSize Uniqueness
    t2))
        f' ScalarTypeBase MonoSize u
    t = ScalarTypeBase MonoSize u
    t

    -- | arrowArg takes a return type and returns it
    -- with the existentials bound moved at the right of arrows.
    -- It also gives the new set of parameters to consider.
    arrowArg ::
      S.Set VName -> -- scope
      S.Set VName -> -- set of argument
      [VName] -> -- size parameters
      RetTypeBase Size as ->
      (RetTypeBase Size as, S.Set VName)
    arrowArg :: forall as.
    Set VName
    -> Set VName
    -> [VName]
    -> RetTypeBase Exp as
    -> (RetTypeBase Exp as, Set VName)
    arrowArg Set VName
    scope Set VName
    argset [VName]
    args_params RetTypeBase Exp as
    rety =
      let (RetTypeBase Exp as
    rety', (Set VName
    funArgs, Set VName
    _)) = Writer (Set VName, Set VName) (RetTypeBase Exp as)
    -> (RetTypeBase Exp as, (Set VName, Set VName))
    forall w a. Writer w a -> (a, w)
    runWriter ((Set VName, [VName])
    -> Set VName
    -> RetTypeBase Exp as
    -> Writer (Set VName, Set VName) (RetTypeBase Exp as)
    forall as'.
    (Set VName, [VName])
    -> Set VName
    -> RetTypeBase Exp as'
    -> Writer (Set VName, Set VName) (RetTypeBase Exp as')
    arrowArgRetType (Set VName
    scope, [VName]
    forall a. Monoid a => a
    mempty) Set VName
    argset RetTypeBase Exp as
    rety)
          new_params :: Set VName
    new_params = Set VName
    funArgs Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList [VName]
    args_params
       in (Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
    forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
    arrowCleanRetType Set VName
    new_params RetTypeBase Exp as
    rety', Set VName
    new_params)
      where
        -- \| takes a type (or return type) and returns it
        -- with the existentials bound moved at the right of arrows.
        -- It also gives (through writer monad) size variables used in arrow arguments
        -- and variables that are constructively used.
        -- The returned type should be cleanned, as too many existentials are introduced.
        arrowArgRetType ::
          (S.Set VName, [VName]) ->
          S.Set VName ->
          RetTypeBase Size as' ->
          Writer (S.Set VName, S.Set VName) (RetTypeBase Size as')
        arrowArgRetType :: forall as'.
    (Set VName, [VName])
    -> Set VName
    -> RetTypeBase Exp as'
    -> Writer (Set VName, Set VName) (RetTypeBase Exp as')
    arrowArgRetType (Set VName
    scope', [VName]
    dimsToPush) Set VName
    argset' (RetType [VName]
    dims TypeBase Exp as'
    ty) = WriterT
      (Set VName, Set VName)
      Identity
      (RetTypeBase Exp as',
       (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp as')
    forall a.
    WriterT
      (Set VName, Set VName)
      Identity
      (a, (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT (Set VName, Set VName) Identity a
    forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
    pass (WriterT
       (Set VName, Set VName)
       Identity
       (RetTypeBase Exp as',
        (Set VName, Set VName) -> (Set VName, Set VName))
     -> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp as'))
    -> WriterT
         (Set VName, Set VName)
         Identity
         (RetTypeBase Exp as',
          (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp as')
    forall a b. (a -> b) -> a -> b
    $ do
          let dims' :: [VName]
    dims' = [VName]
    dims [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    dimsToPush
          (TypeBase Exp as'
    ty', (Set VName
    _, Set VName
    canExt)) <- WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
    -> WriterT
         (Set VName, Set VName)
         Identity
         (TypeBase Exp as', (Set VName, Set VName))
    forall a.
    WriterT (Set VName, Set VName) Identity a
    -> WriterT
         (Set VName, Set VName) Identity (a, (Set VName, Set VName))
    forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
    listen (WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
     -> WriterT
          (Set VName, Set VName)
          Identity
          (TypeBase Exp as', (Set VName, Set VName)))
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
    -> WriterT
         (Set VName, Set VName)
         Identity
         (TypeBase Exp as', (Set VName, Set VName))
    forall a b. (a -> b) -> a -> b
    $ (Set VName, [VName])
    -> TypeBase Exp as'
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
    forall as'.
    (Set VName, [VName])
    -> TypeBase Exp as'
    -> Writer (Set VName, Set VName) (TypeBase Exp as')
    arrowArgType (Set VName
    argset' Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` Set VName
    scope', [VName]
    dims') TypeBase Exp as'
    ty
          (RetTypeBase Exp as',
     (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT
         (Set VName, Set VName)
         Identity
         (RetTypeBase Exp as',
          (Set VName, Set VName) -> (Set VName, Set VName))
    forall a. a -> WriterT (Set VName, Set VName) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([VName] -> TypeBase Exp as' -> RetTypeBase Exp as'
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType ((VName -> Bool) -> [VName] -> [VName]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.member` Set VName
    canExt) [VName]
    dims') TypeBase Exp as'
    ty', (Set VName -> Set VName)
    -> (Set VName, Set VName) -> (Set VName, Set VName)
    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 (Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.difference` Set VName
    canExt))

        arrowArgScalar :: (Set VName, [VName])
    -> ScalarTypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    arrowArgScalar (Set VName, [VName])
    env (Record Map Name (TypeBase Exp u)
    fs) =
          Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    Record (Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u)
    -> WriterT
         (Set VName, Set VName) Identity (Map Name (TypeBase Exp u))
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (TypeBase Exp u
     -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
    -> Map Name (TypeBase Exp u)
    -> WriterT
         (Set VName, Set VName) Identity (Map Name (TypeBase Exp u))
    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) -> Map Name a -> f (Map Name b)
    traverse ((Set VName, [VName])
    -> TypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u)
    forall as'.
    (Set VName, [VName])
    -> TypeBase Exp as'
    -> Writer (Set VName, Set VName) (TypeBase Exp as')
    arrowArgType (Set VName, [VName])
    env) Map Name (TypeBase Exp u)
    fs
        arrowArgScalar (Set VName, [VName])
    env (Sum Map Name [TypeBase Exp u]
    cs) =
          Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u
    forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
    Sum (Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u)
    -> WriterT
         (Set VName, Set VName) Identity (Map Name [TypeBase Exp u])
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (([TypeBase Exp u]
     -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u])
    -> Map Name [TypeBase Exp u]
    -> WriterT
         (Set VName, Set VName) Identity (Map Name [TypeBase Exp u])
    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) -> Map Name a -> f (Map Name b)
    traverse (([TypeBase Exp u]
      -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u])
     -> Map Name [TypeBase Exp u]
     -> WriterT
          (Set VName, Set VName) Identity (Map Name [TypeBase Exp u]))
    -> ((TypeBase Exp u
         -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
        -> [TypeBase Exp u]
        -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u])
    -> (TypeBase Exp u
        -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
    -> Map Name [TypeBase Exp u]
    -> WriterT
         (Set VName, Set VName) Identity (Map Name [TypeBase Exp u])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (TypeBase Exp u
     -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
    -> [TypeBase Exp u]
    -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u]
    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) ((Set VName, [VName])
    -> TypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u)
    forall as'.
    (Set VName, [VName])
    -> TypeBase Exp as'
    -> Writer (Set VName, Set VName) (TypeBase Exp as')
    arrowArgType (Set VName, [VName])
    env) Map Name [TypeBase Exp u]
    cs
        arrowArgScalar (Set VName
    scope', [VName]
    dimsToPush) (Arrow u
    as PName
    argName Diet
    d StructType
    argT RetTypeBase Exp Uniqueness
    retT) =
          WriterT
      (Set VName, Set VName)
      Identity
      (ScalarTypeBase Exp u,
       (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    forall a.
    WriterT
      (Set VName, Set VName)
      Identity
      (a, (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT (Set VName, Set VName) Identity a
    forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
    pass (WriterT
       (Set VName, Set VName)
       Identity
       (ScalarTypeBase Exp u,
        (Set VName, Set VName) -> (Set VName, Set VName))
     -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u))
    -> WriterT
         (Set VName, Set VName)
         Identity
         (ScalarTypeBase Exp u,
          (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    forall a b. (a -> b) -> a -> b
    $ do
            let intros :: Set VName
    intros = (VName -> Bool) -> Set VName -> Set VName
    forall a. (a -> Bool) -> Set a -> Set a
    S.filter VName -> Bool
    notIntrisic Set VName
    argset' Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.difference` Set VName
    scope'
            RetTypeBase Exp Uniqueness
    retT' <- (Set VName, [VName])
    -> Set VName
    -> RetTypeBase Exp Uniqueness
    -> Writer (Set VName, Set VName) (RetTypeBase Exp Uniqueness)
    forall as'.
    (Set VName, [VName])
    -> Set VName
    -> RetTypeBase Exp as'
    -> Writer (Set VName, Set VName) (RetTypeBase Exp as')
    arrowArgRetType (Set VName
    scope', (VName -> Bool) -> [VName] -> [VName]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.notMember` Set VName
    intros) [VName]
    dimsToPush) Set VName
    fullArgset RetTypeBase Exp Uniqueness
    retT
            (ScalarTypeBase Exp u,
     (Set VName, Set VName) -> (Set VName, Set VName))
    -> WriterT
         (Set VName, Set VName)
         Identity
         (ScalarTypeBase Exp u,
          (Set VName, Set VName) -> (Set VName, Set VName))
    forall a. a -> WriterT (Set VName, Set VName) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (u
    -> PName
    -> Diet
    -> StructType
    -> RetTypeBase Exp Uniqueness
    -> ScalarTypeBase Exp u
    forall dim u.
    u
    -> PName
    -> Diet
    -> TypeBase dim NoUniqueness
    -> RetTypeBase dim Uniqueness
    -> ScalarTypeBase dim u
    Arrow u
    as PName
    argName Diet
    d StructType
    argT RetTypeBase Exp Uniqueness
    retT', (Set VName -> Set VName)
    -> (Set VName -> Set VName)
    -> (Set VName, Set VName)
    -> (Set VName, Set VName)
    forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
    forall (p :: * -> * -> *) a b c d.
    Bifunctor p =>
    (a -> b) -> (c -> d) -> p a c -> p b d
    bimap (Set VName
    intros `S.union`) (Set VName -> Set VName -> Set VName
    forall a b. a -> b -> a
    const Set VName
    forall a. Monoid a => a
    mempty))
          where
            notIntrisic :: VName -> Bool
    notIntrisic VName
    vn = VName -> Int
    baseTag VName
    vn Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    maxIntrinsicTag
            argset' :: Set VName
    argset' = FV -> Set VName
    fvVars (FV -> Set VName) -> FV -> Set VName
    forall a b. (a -> b) -> a -> b
    $ StructType -> FV
    forall u. TypeBase Exp u -> FV
    freeInType StructType
    argT
            fullArgset :: Set VName
    fullArgset =
              Set VName
    argset'
                Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> case PName
    argName of
                  PName
    Unnamed -> Set VName
    forall a. Monoid a => a
    mempty
                  Named VName
    vn -> VName -> Set VName
    forall a. a -> Set a
    S.singleton VName
    vn
        arrowArgScalar (Set VName, [VName])
    env (TypeVar u
    u QualName VName
    qn [TypeArg Exp]
    args) =
          u -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp u
    forall dim u.
    u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
    TypeVar u
    u QualName VName
    qn ([TypeArg Exp] -> ScalarTypeBase Exp u)
    -> WriterT (Set VName, Set VName) Identity [TypeArg Exp]
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (TypeArg Exp
     -> WriterT (Set VName, Set VName) Identity (TypeArg Exp))
    -> [TypeArg Exp]
    -> WriterT (Set VName, Set VName) Identity [TypeArg 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 TypeArg Exp
    -> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
    arrowArgArg [TypeArg Exp]
    args
          where
            arrowArgArg :: TypeArg Exp
    -> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
    arrowArgArg (TypeArgDim Exp
    dim) = Exp -> TypeArg Exp
    forall dim. dim -> TypeArg dim
    TypeArgDim (Exp -> TypeArg Exp)
    -> WriterT (Set VName, Set VName) Identity Exp
    -> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> WriterT (Set VName, Set VName) Identity Exp
    forall {a} {a} {m :: * -> *} {f :: * -> *}.
    (MonadWriter (a, Set a) m, Monoid a, Ord a) =>
    ExpBase f a -> m (ExpBase f a)
    arrowArgSize Exp
    dim
            arrowArgArg (TypeArgType StructType
    ty) = StructType -> TypeArg Exp
    forall dim. TypeBase dim NoUniqueness -> TypeArg dim
    TypeArgType (StructType -> TypeArg Exp)
    -> WriterT (Set VName, Set VName) Identity StructType
    -> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Set VName, [VName])
    -> StructType -> WriterT (Set VName, Set VName) Identity StructType
    forall as'.
    (Set VName, [VName])
    -> TypeBase Exp as'
    -> Writer (Set VName, Set VName) (TypeBase Exp as')
    arrowArgType (Set VName, [VName])
    env StructType
    ty
        arrowArgScalar (Set VName, [VName])
    _ ScalarTypeBase Exp u
    ty = ScalarTypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    forall a. a -> WriterT (Set VName, Set VName) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ScalarTypeBase Exp u
    ty

        arrowArgType ::
          (S.Set VName, [VName]) ->
          TypeBase Size as' ->
          Writer (S.Set VName, S.Set VName) (TypeBase Size as')
        arrowArgType :: forall as'.
    (Set VName, [VName])
    -> TypeBase Exp as'
    -> Writer (Set VName, Set VName) (TypeBase Exp as')
    arrowArgType (Set VName, [VName])
    env (Array as'
    u Shape Exp
    shape ScalarTypeBase Exp NoUniqueness
    scalar) =
          as'
    -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as'
    forall dim u.
    u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
    Array as'
    u (Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as')
    -> WriterT (Set VName, Set VName) Identity (Shape Exp)
    -> WriterT
         (Set VName, Set VName)
         Identity
         (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as')
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> WriterT (Set VName, Set VName) Identity Exp)
    -> Shape Exp -> WriterT (Set VName, Set VName) Identity (Shape Exp)
    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) -> Shape a -> f (Shape b)
    traverse Exp -> WriterT (Set VName, Set VName) Identity Exp
    forall {a} {a} {m :: * -> *} {f :: * -> *}.
    (MonadWriter (a, Set a) m, Monoid a, Ord a) =>
    ExpBase f a -> m (ExpBase f a)
    arrowArgSize Shape Exp
    shape WriterT
      (Set VName, Set VName)
      Identity
      (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as')
    -> WriterT
         (Set VName, Set VName) Identity (ScalarTypeBase Exp NoUniqueness)
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
    forall a b.
    WriterT (Set VName, Set VName) Identity (a -> b)
    -> WriterT (Set VName, Set VName) Identity a
    -> WriterT (Set VName, Set VName) Identity b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (Set VName, [VName])
    -> ScalarTypeBase Exp NoUniqueness
    -> WriterT
         (Set VName, Set VName) Identity (ScalarTypeBase Exp NoUniqueness)
    forall {u}.
    (Set VName, [VName])
    -> ScalarTypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    arrowArgScalar (Set VName, [VName])
    env ScalarTypeBase Exp NoUniqueness
    scalar
        arrowArgType (Set VName, [VName])
    env (Scalar ScalarTypeBase Exp as'
    ty) =
          ScalarTypeBase Exp as' -> TypeBase Exp as'
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase Exp as' -> TypeBase Exp as')
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp as')
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Set VName, [VName])
    -> ScalarTypeBase Exp as'
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp as')
    forall {u}.
    (Set VName, [VName])
    -> ScalarTypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
    arrowArgScalar (Set VName, [VName])
    env ScalarTypeBase Exp as'
    ty

        arrowArgSize :: ExpBase f a -> m (ExpBase f a)
    arrowArgSize s :: ExpBase f a
    s@(Var QualName a
    qn f StructType
    _ SrcLoc
    _) = (ExpBase f a, (a, Set a)) -> m (ExpBase f a)
    forall a. (a, (a, Set a)) -> m a
    forall w (m :: * -> *) a. MonadWriter w m => (a, w) -> m a
    writer (ExpBase f a
    s, (a
    forall a. Monoid a => a
    mempty, a -> Set a
    forall a. a -> Set a
    S.singleton (a -> Set a) -> a -> Set a
    forall a b. (a -> b) -> a -> b
    $ QualName a -> a
    forall vn. QualName vn -> vn
    qualLeaf QualName a
    qn))
        arrowArgSize ExpBase f a
    s = ExpBase f a -> m (ExpBase f a)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ExpBase f a
    s

        -- \| arrowClean cleans the mess in the type
        arrowCleanRetType :: S.Set VName -> RetTypeBase Size as -> RetTypeBase Size as
        arrowCleanRetType :: forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
    arrowCleanRetType Set VName
    paramed (RetType [VName]
    dims TypeBase Exp as
    ty) =
          [VName] -> TypeBase Exp as -> RetTypeBase Exp as
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType ([VName] -> [VName]
    forall a. Ord a => [a] -> [a]
    nubOrd ([VName] -> [VName]) -> [VName] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ (VName -> Bool) -> [VName] -> [VName]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.notMember` Set VName
    paramed) [VName]
    dims) (Set VName -> TypeBase Exp as -> TypeBase Exp as
    forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
    arrowCleanType (Set VName
    paramed Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    `S.union` [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList [VName]
    dims) TypeBase Exp as
    ty)

        arrowCleanScalar :: S.Set VName -> ScalarTypeBase Size as -> ScalarTypeBase Size as
        arrowCleanScalar :: forall as.
    Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
    arrowCleanScalar Set VName
    paramed (Record Map Name (TypeBase Exp as)
    fs) =
          Map Name (TypeBase Exp as) -> ScalarTypeBase Exp as
    forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
    Record (Map Name (TypeBase Exp as) -> ScalarTypeBase Exp as)
    -> Map Name (TypeBase Exp as) -> ScalarTypeBase Exp as
    forall a b. (a -> b) -> a -> b
    $ (TypeBase Exp as -> TypeBase Exp as)
    -> Map Name (TypeBase Exp as) -> Map Name (TypeBase Exp as)
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map (Set VName -> TypeBase Exp as -> TypeBase Exp as
    forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
    arrowCleanType Set VName
    paramed) Map Name (TypeBase Exp as)
    fs
        arrowCleanScalar Set VName
    paramed (Sum Map Name [TypeBase Exp as]
    cs) =
          Map Name [TypeBase Exp as] -> ScalarTypeBase Exp as
    forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
    Sum (Map Name [TypeBase Exp as] -> ScalarTypeBase Exp as)
    -> Map Name [TypeBase Exp as] -> ScalarTypeBase Exp as
    forall a b. (a -> b) -> a -> b
    $ (([TypeBase Exp as] -> [TypeBase Exp as])
    -> Map Name [TypeBase Exp as] -> Map Name [TypeBase Exp as]
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map (([TypeBase Exp as] -> [TypeBase Exp as])
     -> Map Name [TypeBase Exp as] -> Map Name [TypeBase Exp as])
    -> ((TypeBase Exp as -> TypeBase Exp as)
        -> [TypeBase Exp as] -> [TypeBase Exp as])
    -> (TypeBase Exp as -> TypeBase Exp as)
    -> Map Name [TypeBase Exp as]
    -> Map Name [TypeBase Exp as]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (TypeBase Exp as -> TypeBase Exp as)
    -> [TypeBase Exp as] -> [TypeBase Exp as]
    forall a b. (a -> b) -> [a] -> [b]
    map) (Set VName -> TypeBase Exp as -> TypeBase Exp as
    forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
    arrowCleanType Set VName
    paramed) Map Name [TypeBase Exp as]
    cs
        arrowCleanScalar Set VName
    paramed (Arrow as
    as PName
    argName Diet
    d StructType
    argT RetTypeBase Exp Uniqueness
    retT) =
          as
    -> PName
    -> Diet
    -> StructType
    -> RetTypeBase Exp Uniqueness
    -> ScalarTypeBase Exp as
    forall dim u.
    u
    -> PName
    -> Diet
    -> TypeBase dim NoUniqueness
    -> RetTypeBase dim Uniqueness
    -> ScalarTypeBase dim u
    Arrow as
    as PName
    argName Diet
    d StructType
    argT (Set VName
    -> RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
    arrowCleanRetType Set VName
    paramed RetTypeBase Exp Uniqueness
    retT)
        arrowCleanScalar Set VName
    paramed (TypeVar as
    u QualName VName
    qn [TypeArg Exp]
    args) =
          as -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp as
    forall dim u.
    u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
    TypeVar as
    u QualName VName
    qn ([TypeArg Exp] -> ScalarTypeBase Exp as)
    -> [TypeArg Exp] -> ScalarTypeBase Exp as
    forall a b. (a -> b) -> a -> b
    $ (TypeArg Exp -> TypeArg Exp) -> [TypeArg Exp] -> [TypeArg Exp]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeArg Exp -> TypeArg Exp
    arrowCleanArg [TypeArg Exp]
    args
          where
            arrowCleanArg :: TypeArg Exp -> TypeArg Exp
    arrowCleanArg (TypeArgDim Exp
    dim) = Exp -> TypeArg Exp
    forall dim. dim -> TypeArg dim
    TypeArgDim Exp
    dim
            arrowCleanArg (TypeArgType StructType
    ty) = StructType -> TypeArg Exp
    forall dim. TypeBase dim NoUniqueness -> TypeArg dim
    TypeArgType (StructType -> TypeArg Exp) -> StructType -> TypeArg Exp
    forall a b. (a -> b) -> a -> b
    $ Set VName -> StructType -> StructType
    forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
    arrowCleanType Set VName
    paramed StructType
    ty
        arrowCleanScalar Set VName
    _ ScalarTypeBase Exp as
    ty = ScalarTypeBase Exp as
    ty

        arrowCleanType :: S.Set VName -> TypeBase Size as -> TypeBase Size as
        arrowCleanType :: forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
    arrowCleanType Set VName
    paramed (Array as
    u Shape Exp
    shape ScalarTypeBase Exp NoUniqueness
    scalar) =
          as
    -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as
    forall dim u.
    u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
    Array as
    u Shape Exp
    shape (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as)
    -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as
    forall a b. (a -> b) -> a -> b
    $ Set VName
    -> ScalarTypeBase Exp NoUniqueness
    -> ScalarTypeBase Exp NoUniqueness
    forall as.
    Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
    arrowCleanScalar Set VName
    paramed ScalarTypeBase Exp NoUniqueness
    scalar
        arrowCleanType Set VName
    paramed (Scalar ScalarTypeBase Exp as
    ty) =
          ScalarTypeBase Exp as -> TypeBase Exp as
    forall dim u. ScalarTypeBase dim u -> TypeBase dim u
    Scalar (ScalarTypeBase Exp as -> TypeBase Exp as)
    -> ScalarTypeBase Exp as -> TypeBase Exp as
    forall a b. (a -> b) -> a -> b
    $ Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
    forall as.
    Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
    arrowCleanScalar Set VName
    paramed ScalarTypeBase Exp as
    ty

    -- Monomorphise a polymorphic function at the types given in the instance
    -- list. Monomorphises the body of the function as well. Returns the fresh name
    -- of the generated monomorphic function and its 'ValBind' representation.
    monomorphiseBinding ::
      PolyBinding ->
      MonoType ->
      MonoM (VName, InferSizeArgs, ValBind)
    monomorphiseBinding :: PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
    monomorphiseBinding (PolyBinding (Maybe EntryPoint
    entry, VName
    name, [TypeParam]
    tparams, [Pat ParamType]
    params, RetTypeBase Exp Uniqueness
    rettype, Exp
    body, [AttrInfo VName]
    attrs, SrcLoc
    loc)) MonoType
    inst_t = MonoM (VName, InferSizeArgs, ValBind)
    -> MonoM (VName, InferSizeArgs, ValBind)
    forall a. MonoM a -> MonoM a
    isolateNormalisation (MonoM (VName, InferSizeArgs, ValBind)
     -> MonoM (VName, InferSizeArgs, ValBind))
    -> MonoM (VName, InferSizeArgs, ValBind)
    -> MonoM (VName, InferSizeArgs, ValBind)
    forall a b. (a -> b) -> a -> b
    $ do
      let bind_t :: StructType
    bind_t = [Pat ParamType] -> RetTypeBase Exp Uniqueness -> StructType
    funType [Pat ParamType]
    params RetTypeBase Exp Uniqueness
    rettype
      (Map VName StructRetType
    substs, [TypeParam]
    t_shape_params) <-
        SrcLoc
    -> StructType
    -> MonoType
    -> MonoM (Map VName StructRetType, [TypeParam])
    forall (m :: * -> *).
    MonadFreshNames m =>
    SrcLoc
    -> StructType
    -> MonoType
    -> m (Map VName StructRetType, [TypeParam])
    typeSubstsM SrcLoc
    loc StructType
    bind_t (MonoType -> MonoM (Map VName StructRetType, [TypeParam]))
    -> MonoType -> MonoM (Map VName StructRetType, [TypeParam])
    forall a b. (a -> b) -> a -> b
    $ MonoType -> MonoType
    noNamedParams MonoType
    inst_t
      let shape_names :: Set VName
    shape_names = [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
    forall a b. (a -> b) -> a -> b
    $ (TypeParam -> VName) -> [TypeParam] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeParam -> VName
    forall vn. TypeParamBase vn -> vn
    typeParamName ([TypeParam] -> [VName]) -> [TypeParam] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ [TypeParam]
    shape_params [TypeParam] -> [TypeParam] -> [TypeParam]
    forall a. [a] -> [a] -> [a]
    ++ [TypeParam]
    t_shape_params
          substs' :: Map VName (Subst StructRetType)
    substs' = (StructRetType -> Subst StructRetType)
    -> Map VName StructRetType -> Map VName (Subst StructRetType)
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map ([TypeParam] -> StructRetType -> Subst StructRetType
    forall t. [TypeParam] -> t -> Subst t
    Subst []) Map VName StructRetType
    substs
          substStructType :: ParamType -> ParamType
    substStructType =
            (VName -> Maybe (Subst (RetTypeBase Exp Diet)))
    -> ParamType -> ParamType
    forall u.
    Monoid u =>
    (VName -> Maybe (Subst (RetTypeBase Exp u)))
    -> TypeBase Exp u -> TypeBase Exp u
    substTypesAny ((Subst StructRetType -> Subst (RetTypeBase Exp Diet))
    -> Maybe (Subst StructRetType)
    -> Maybe (Subst (RetTypeBase Exp Diet))
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((StructRetType -> RetTypeBase Exp Diet)
    -> Subst StructRetType -> Subst (RetTypeBase Exp Diet)
    forall a b. (a -> b) -> Subst a -> Subst b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((NoUniqueness -> Diet) -> StructRetType -> RetTypeBase Exp Diet
    forall b c a. (b -> c) -> RetTypeBase a b -> RetTypeBase a c
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (Diet -> NoUniqueness -> Diet
    forall a b. a -> b -> a
    const Diet
    forall a. Monoid a => a
    mempty))) (Maybe (Subst StructRetType)
     -> Maybe (Subst (RetTypeBase Exp Diet)))
    -> TypeSubs -> VName -> Maybe (Subst (RetTypeBase Exp Diet))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName
    -> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName (Subst StructRetType)
    substs'))
          params' :: [Pat ParamType]
    params' = (Pat ParamType -> Pat ParamType)
    -> [Pat ParamType] -> [Pat ParamType]
    forall a b. (a -> b) -> [a] -> [b]
    map ((ParamType -> ParamType) -> Pat ParamType -> Pat ParamType
    forall t. (t -> t) -> Pat t -> Pat t
    substPat ParamType -> ParamType
    substStructType) [Pat ParamType]
    params
      [Pat ParamType]
    params'' <- Set VName -> MonoM [Pat ParamType] -> MonoM [Pat ParamType]
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs Set VName
    shape_names (MonoM [Pat ParamType] -> MonoM [Pat ParamType])
    -> MonoM [Pat ParamType] -> MonoM [Pat ParamType]
    forall a b. (a -> b) -> a -> b
    $ (Pat ParamType -> MonoM (Pat ParamType))
    -> [Pat ParamType] -> MonoM [Pat ParamType]
    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 Pat ParamType -> MonoM (Pat ParamType)
    forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
    transformPat [Pat ParamType]
    params'
      ExpReplacements
    exp_naming <- MonoM ExpReplacements
    forall s (m :: * -> *). MonadState s m => m s
    get MonoM ExpReplacements -> MonoM () -> MonoM ExpReplacements
    forall a b. MonoM a -> MonoM b -> MonoM a
    forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
    <* ExpReplacements -> MonoM ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put ExpReplacements
    forall a. Monoid a => a
    mempty

      let args :: Set VName
    args = [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
    forall a b. (a -> b) -> a -> b
    $ (Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Pat ParamType -> [VName]
    forall t. Pat t -> [VName]
    patNames [Pat ParamType]
    params
          arg_params :: [VName]
    arg_params = ((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (ReplacedExp, VName) -> VName
    forall a b. (a, b) -> b
    snd ExpReplacements
    exp_naming

      RetTypeBase Exp Uniqueness
    rettype' <-
        ExpReplacements
    -> MonoM (RetTypeBase Exp Uniqueness)
    -> MonoM (RetTypeBase Exp Uniqueness)
    forall a. ExpReplacements -> MonoM a -> MonoM a
    withParams ExpReplacements
    exp_naming (MonoM (RetTypeBase Exp Uniqueness)
     -> MonoM (RetTypeBase Exp Uniqueness))
    -> MonoM (RetTypeBase Exp Uniqueness)
    -> MonoM (RetTypeBase Exp Uniqueness)
    forall a b. (a -> b) -> a -> b
    $
          Set VName
    -> MonoM (RetTypeBase Exp Uniqueness)
    -> MonoM (RetTypeBase Exp Uniqueness)
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs (Set VName
    args Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> Set VName
    shape_names) (MonoM (RetTypeBase Exp Uniqueness)
     -> MonoM (RetTypeBase Exp Uniqueness))
    -> MonoM (RetTypeBase Exp Uniqueness)
    -> MonoM (RetTypeBase Exp Uniqueness)
    forall a b. (a -> b) -> a -> b
    $
            RetTypeBase Exp Uniqueness -> MonoM (RetTypeBase Exp Uniqueness)
    forall {as}. RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    hardTransformRetType (TypeSubs
    -> RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall a. Substitutable a => TypeSubs -> a -> a
    applySubst (VName
    -> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName (Subst StructRetType)
    substs') RetTypeBase Exp Uniqueness
    rettype)
      ExpReplacements
    extNaming <- MonoM ExpReplacements
    forall s (m :: * -> *). MonadState s m => m s
    get MonoM ExpReplacements -> MonoM () -> MonoM ExpReplacements
    forall a b. MonoM a -> MonoM b -> MonoM a
    forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
    <* ExpReplacements -> MonoM ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put ExpReplacements
    forall a. Monoid a => a
    mempty
      Set VName
    scope <- Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    S.union Set VName
    shape_names (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM (Set VName)
    askScope'
      let (RetTypeBase Exp Uniqueness
    rettype'', Set VName
    new_params) = Set VName
    -> Set VName
    -> [VName]
    -> RetTypeBase Exp Uniqueness
    -> (RetTypeBase Exp Uniqueness, Set VName)
    forall as.
    Set VName
    -> Set VName
    -> [VName]
    -> RetTypeBase Exp as
    -> (RetTypeBase Exp as, Set VName)
    arrowArg Set VName
    scope Set VName
    args [VName]
    arg_params RetTypeBase Exp Uniqueness
    rettype'
          bind_t' :: StructType
    bind_t' = TypeSubs -> StructType -> StructType
    forall u.
    Monoid u =>
    (VName -> Maybe (Subst (RetTypeBase Exp u)))
    -> TypeBase Exp u -> TypeBase Exp u
    substTypesAny (VName
    -> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName (Subst StructRetType)
    substs') StructType
    bind_t
          mkExplicit :: VName -> Bool
    mkExplicit =
            (VName -> Set VName -> Bool) -> Set VName -> VName -> Bool
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip
              VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    S.member
              (StructType -> Set VName
    mustBeExplicitInBinding StructType
    bind_t'' Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> StructType -> Set VName
    mustBeExplicitInBinding StructType
    bind_t')
          ([TypeParam]
    shape_params_explicit, [TypeParam]
    shape_params_implicit) =
            (TypeParam -> Bool) -> [TypeParam] -> ([TypeParam], [TypeParam])
    forall a. (a -> Bool) -> [a] -> ([a], [a])
    partition (VName -> Bool
    mkExplicit (VName -> Bool) -> (TypeParam -> VName) -> TypeParam -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeParam -> VName
    forall vn. TypeParamBase vn -> vn
    typeParamName) ([TypeParam] -> ([TypeParam], [TypeParam]))
    -> [TypeParam] -> ([TypeParam], [TypeParam])
    forall a b. (a -> b) -> a -> b
    $
              [TypeParam]
    shape_params [TypeParam] -> [TypeParam] -> [TypeParam]
    forall a. [a] -> [a] -> [a]
    ++ [TypeParam]
    t_shape_params [TypeParam] -> [TypeParam] -> [TypeParam]
    forall a. [a] -> [a] -> [a]
    ++ (VName -> TypeParam) -> [VName] -> [TypeParam]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> SrcLoc -> TypeParam
    forall vn. vn -> SrcLoc -> TypeParamBase vn
    `TypeParamDim` SrcLoc
    forall a. Monoid a => a
    mempty) (Set VName -> [VName]
    forall a. Set a -> [a]
    S.toList Set VName
    new_params)
          exp_naming' :: ExpReplacements
    exp_naming' = ((ReplacedExp, VName) -> Bool)
    -> ExpReplacements -> ExpReplacements
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.member` Set VName
    new_params) (VName -> Bool)
    -> ((ReplacedExp, VName) -> VName) -> (ReplacedExp, VName) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (ReplacedExp, VName) -> VName
    forall a b. (a, b) -> b
    snd) (ExpReplacements
    extNaming ExpReplacements -> ExpReplacements -> ExpReplacements
    forall a. Semigroup a => a -> a -> a
    <> ExpReplacements
    exp_naming)

          bind_t'' :: StructType
    bind_t'' = [Pat ParamType] -> RetTypeBase Exp Uniqueness -> StructType
    funType [Pat ParamType]
    params'' RetTypeBase Exp Uniqueness
    rettype''
          bind_r :: ExpReplacements
    bind_r = ExpReplacements
    exp_naming ExpReplacements -> ExpReplacements -> ExpReplacements
    forall a. Semigroup a => a -> a -> a
    <> ExpReplacements
    extNaming
      Exp
    body' <- TypeSubs -> Exp -> MonoM Exp
    forall {m :: * -> *}. Monad m => TypeSubs -> Exp -> m Exp
    updateExpTypes (VName
    -> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName (Subst StructRetType)
    substs') Exp
    body
      Exp
    body'' <- ExpReplacements -> MonoM Exp -> MonoM Exp
    forall a. ExpReplacements -> MonoM a -> MonoM a
    withParams ExpReplacements
    exp_naming' (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Set VName -> MonoM Exp -> MonoM Exp
    forall a. Set VName -> MonoM a -> MonoM a
    withArgs (Set VName
    shape_names Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> Set VName
    args) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
    forall a b. (a -> b) -> a -> b
    $ Exp -> MonoM Exp
    transformExp Exp
    body'
      Set VName
    scope' <- Set VName -> Set VName -> Set VName
    forall a. Ord a => Set a -> Set a -> Set a
    S.union (Set VName
    shape_names Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> Set VName
    args) (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM (Set VName)
    askScope'
      Exp
    body''' <-
        ExpReplacements -> Exp -> Exp
    expReplace ExpReplacements
    exp_naming' (Exp -> Exp) -> MonoM Exp -> MonoM Exp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Exp -> ExpReplacements -> MonoM Exp
    calculateDims Exp
    body'' (ExpReplacements -> MonoM Exp)
    -> (ExpReplacements -> ExpReplacements)
    -> ExpReplacements
    -> MonoM Exp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Set VName -> ExpReplacements -> ExpReplacements
    canCalculate Set VName
    scope' (ExpReplacements -> MonoM Exp)
    -> MonoM ExpReplacements -> MonoM Exp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< MonoM ExpReplacements
    forall s (m :: * -> *). MonadState s m => m s
    get)

      Bool
    seen_before <- VName -> [VName] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem VName
    name ([VName] -> Bool) -> (Lifts -> [VName]) -> Lifts -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
    -> Lifts -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName, MonoType) -> VName
    forall a b. (a, b) -> a
    fst ((VName, MonoType) -> VName)
    -> (((VName, MonoType), (VName, InferSizeArgs))
        -> (VName, MonoType))
    -> ((VName, MonoType), (VName, InferSizeArgs))
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, MonoType), (VName, InferSizeArgs)) -> (VName, MonoType)
    forall a b. (a, b) -> a
    fst) (Lifts -> Bool) -> MonoM Lifts -> MonoM Bool
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM Lifts
    getLifts
      VName
    name' <-
        if [TypeParam] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [TypeParam]
    tparams Bool -> Bool -> Bool
    && Maybe EntryPoint -> Bool
    forall a. Maybe a -> Bool
    isNothing Maybe EntryPoint
    entry Bool -> Bool -> Bool
    && Bool -> Bool
    not Bool
    seen_before
          then VName -> MonoM VName
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    name
          else VName -> MonoM VName
    forall (m :: * -> *). MonadFreshNames m => VName -> m VName
    newName VName
    name

      (VName, InferSizeArgs, ValBind)
    -> MonoM (VName, InferSizeArgs, ValBind)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
        ( VName
    name',
          -- If the function is an entry point, then it cannot possibly
          -- need any explicit size arguments (checked by type checker).
          if Maybe EntryPoint -> Bool
    forall a. Maybe a -> Bool
    isJust Maybe EntryPoint
    entry
            then MonoM [Exp] -> InferSizeArgs
    forall a b. a -> b -> a
    const (MonoM [Exp] -> InferSizeArgs) -> MonoM [Exp] -> InferSizeArgs
    forall a b. (a -> b) -> a -> b
    $ [Exp] -> MonoM [Exp]
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure []
            else [TypeParam] -> StructType -> ExpReplacements -> InferSizeArgs
    inferSizeArgs [TypeParam]
    shape_params_explicit StructType
    bind_t'' ExpReplacements
    bind_r,
          if Maybe EntryPoint -> Bool
    forall a. Maybe a -> Bool
    isJust Maybe EntryPoint
    entry
            then
              VName
    -> [TypeParam]
    -> [Pat ParamType]
    -> RetTypeBase Exp Uniqueness
    -> Exp
    -> ValBind
    toValBinding
                VName
    name'
                ([TypeParam]
    shape_params_explicit [TypeParam] -> [TypeParam] -> [TypeParam]
    forall a. [a] -> [a] -> [a]
    ++ [TypeParam]
    shape_params_implicit)
                [Pat ParamType]
    params''
                RetTypeBase Exp Uniqueness
    rettype''
                (ExpReplacements -> Exp -> Exp
    entryAssert ExpReplacements
    exp_naming Exp
    body''')
            else
              VName
    -> [TypeParam]
    -> [Pat ParamType]
    -> RetTypeBase Exp Uniqueness
    -> Exp
    -> ValBind
    toValBinding
                VName
    name'
                [TypeParam]
    shape_params_implicit
                ((TypeParam -> Pat ParamType) -> [TypeParam] -> [Pat ParamType]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeParam -> Pat ParamType
    forall {vn} {dim} {als}.
    TypeParamBase vn -> PatBase Info vn (TypeBase dim als)
    shapeParam [TypeParam]
    shape_params_explicit [Pat ParamType] -> [Pat ParamType] -> [Pat ParamType]
    forall a. [a] -> [a] -> [a]
    ++ [Pat ParamType]
    params'')
                RetTypeBase Exp Uniqueness
    rettype''
                Exp
    body'''
        )
      where
        askScope' :: MonoM (Set VName)
    askScope' = (VName -> Bool) -> Set VName -> Set VName
    forall a. (a -> Bool) -> Set a -> Set a
    S.filter (VName -> [VName] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` RetTypeBase Exp Uniqueness -> [VName]
    forall dim as. RetTypeBase dim as -> [VName]
    retDims RetTypeBase Exp Uniqueness
    rettype) (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> MonoM (Set VName)
    askScope

        shape_params :: [TypeParam]
    shape_params = (TypeParam -> Bool) -> [TypeParam] -> [TypeParam]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Bool -> Bool
    not (Bool -> Bool) -> (TypeParam -> Bool) -> TypeParam -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeParam -> Bool
    forall vn. TypeParamBase vn -> Bool
    isTypeParam) [TypeParam]
    tparams

        updateExpTypes :: TypeSubs -> Exp -> m Exp
    updateExpTypes TypeSubs
    substs = ASTMapper m -> Exp -> m Exp
    forall x (m :: * -> *).
    (ASTMappable x, Monad m) =>
    ASTMapper m -> x -> m x
    forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
    astMap (TypeSubs -> ASTMapper m
    mapper TypeSubs
    substs)

        hardTransformRetType :: RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    hardTransformRetType (RetType [VName]
    dims TypeBase Exp as
    ty) = do
          TypeBase Exp as
    ty' <- TypeBase Exp as -> MonoM (TypeBase Exp as)
    forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
    transformType TypeBase Exp as
    ty
          Set VName
    unbounded <- Set VName -> MonoM (Set VName)
    askIntros (Set VName -> MonoM (Set VName)) -> Set VName -> MonoM (Set VName)
    forall a b. (a -> b) -> a -> b
    $ FV -> Set VName
    fvVars (FV -> Set VName) -> FV -> Set VName
    forall a b. (a -> b) -> a -> b
    $ TypeBase Exp as -> FV
    forall u. TypeBase Exp u -> FV
    freeInType TypeBase Exp as
    ty'
          let dims' :: [VName]
    dims' = Set VName -> [VName]
    forall a. Set a -> [a]
    S.toList Set VName
    unbounded
          RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (RetTypeBase Exp as -> MonoM (RetTypeBase Exp as))
    -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
    forall a b. (a -> b) -> a -> b
    $ [VName] -> TypeBase Exp as -> RetTypeBase Exp as
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType ([VName]
    dims' [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    dims) TypeBase Exp as
    ty'

        mapper :: TypeSubs -> ASTMapper m
    mapper TypeSubs
    substs =
          ASTMapper
            { mapOnExp :: Exp -> m Exp
    mapOnExp = TypeSubs -> Exp -> m Exp
    updateExpTypes TypeSubs
    substs,
              mapOnName :: QualName VName -> m (QualName VName)
    mapOnName = QualName VName -> m (QualName VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure,
              mapOnStructType :: StructType -> m StructType
    mapOnStructType = StructType -> m StructType
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (StructType -> m StructType)
    -> (StructType -> StructType) -> StructType -> m StructType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeSubs -> StructType -> StructType
    forall a. Substitutable a => TypeSubs -> a -> a
    applySubst TypeSubs
    substs,
              mapOnParamType :: ParamType -> m ParamType
    mapOnParamType = ParamType -> m ParamType
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (ParamType -> m ParamType)
    -> (ParamType -> ParamType) -> ParamType -> m ParamType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeSubs -> ParamType -> ParamType
    forall a. Substitutable a => TypeSubs -> a -> a
    applySubst TypeSubs
    substs,
              mapOnResRetType :: RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness)
    mapOnResRetType = RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness))
    -> (RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness)
    -> RetTypeBase Exp Uniqueness
    -> m (RetTypeBase Exp Uniqueness)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeSubs
    -> RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
    forall a. Substitutable a => TypeSubs -> a -> a
    applySubst TypeSubs
    substs
            }

        shapeParam :: TypeParamBase vn -> PatBase Info vn (TypeBase dim als)
    shapeParam TypeParamBase vn
    tp = vn
    -> Info (TypeBase dim als)
    -> SrcLoc
    -> PatBase Info vn (TypeBase dim als)
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id (TypeParamBase vn -> vn
    forall vn. TypeParamBase vn -> vn
    typeParamName TypeParamBase vn
    tp) (TypeBase dim als -> Info (TypeBase dim als)
    forall a. a -> Info a
    Info TypeBase dim als
    forall dim als. TypeBase dim als
    i64) (SrcLoc -> PatBase Info vn (TypeBase dim als))
    -> SrcLoc -> PatBase Info vn (TypeBase dim als)
    forall a b. (a -> b) -> a -> b
    $ TypeParamBase vn -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf TypeParamBase vn
    tp

        toValBinding :: VName
    -> [TypeParam]
    -> [Pat ParamType]
    -> RetTypeBase Exp Uniqueness
    -> Exp
    -> ValBind
    toValBinding VName
    name' [TypeParam]
    tparams' [Pat ParamType]
    params'' RetTypeBase Exp Uniqueness
    rettype' Exp
    body'' =
          ValBind
            { valBindEntryPoint :: Maybe (Info EntryPoint)
    valBindEntryPoint = EntryPoint -> Info EntryPoint
    forall a. a -> Info a
    Info (EntryPoint -> Info EntryPoint)
    -> Maybe EntryPoint -> Maybe (Info EntryPoint)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Maybe EntryPoint
    entry,
              valBindName :: VName
    valBindName = VName
    name',
              valBindRetType :: Info (RetTypeBase Exp Uniqueness)
    valBindRetType = RetTypeBase Exp Uniqueness -> Info (RetTypeBase Exp Uniqueness)
    forall a. a -> Info a
    Info RetTypeBase Exp Uniqueness
    rettype',
              valBindRetDecl :: Maybe (TypeExp Exp VName)
    valBindRetDecl = Maybe (TypeExp Exp VName)
    forall a. Maybe a
    Nothing,
              valBindTypeParams :: [TypeParam]
    valBindTypeParams = [TypeParam]
    tparams',
              valBindParams :: [Pat ParamType]
    valBindParams = [Pat ParamType]
    params'',
              valBindBody :: Exp
    valBindBody = Exp
    body'',
              valBindDoc :: Maybe DocComment
    valBindDoc = Maybe DocComment
    forall a. Maybe a
    Nothing,
              valBindAttrs :: [AttrInfo VName]
    valBindAttrs = [AttrInfo VName]
    attrs,
              valBindLocation :: SrcLoc
    valBindLocation = SrcLoc
    loc
            }

    typeSubstsM ::
      (MonadFreshNames m) =>
      SrcLoc ->
      StructType ->
      MonoType ->
      m (M.Map VName StructRetType, [TypeParam])
    typeSubstsM :: forall (m :: * -> *).
    MonadFreshNames m =>
    SrcLoc
    -> StructType
    -> MonoType
    -> m (Map VName StructRetType, [TypeParam])
    typeSubstsM SrcLoc
    loc StructType
    orig_t1 MonoType
    orig_t2 =
      WriterT [TypeParam] m (Map VName StructRetType)
    -> m (Map VName StructRetType, [TypeParam])
    forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
    runWriterT (WriterT [TypeParam] m (Map VName StructRetType)
     -> m (Map VName StructRetType, [TypeParam]))
    -> WriterT [TypeParam] m (Map VName StructRetType)
    -> m (Map VName StructRetType, [TypeParam])
    forall a b. (a -> b) -> a -> b
    $ (Map VName StructRetType, Map Int Exp) -> Map VName StructRetType
    forall a b. (a, b) -> a
    fst ((Map VName StructRetType, Map Int Exp) -> Map VName StructRetType)
    -> WriterT [TypeParam] m (Map VName StructRetType, Map Int Exp)
    -> WriterT [TypeParam] m (Map VName StructRetType)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StateT
      (Map VName StructRetType, Map Int Exp) (WriterT [TypeParam] m) ()
    -> (Map VName StructRetType, Map Int Exp)
    -> WriterT [TypeParam] m (Map VName StructRetType, Map Int Exp)
    forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
    execStateT (StructType
    -> MonoType
    -> StateT
         (Map VName StructRetType, Map Int Exp) (WriterT [TypeParam] m) ()
    forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
           {m :: * -> *}.
    (MonadWriter [TypeParam] (t (t m)), MonadFreshNames m,
     MonadTrans t, MonadTrans t,
     MonadState (Map VName StructRetType, Map Int Exp) (t (t m)),
     Monad (t m)) =>
    StructType -> MonoType -> t (t m) ()
    sub StructType
    orig_t1 MonoType
    orig_t2) (Map VName StructRetType
    forall a. Monoid a => a
    mempty, Map Int Exp
    forall a. Monoid a => a
    mempty)
      where
        subRet :: StructType -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
    subRet (Scalar (TypeVar NoUniqueness
    _ QualName VName
    v [TypeArg Exp]
    _)) RetTypeBase MonoSize NoUniqueness
    rt =
          Bool -> t (t m) () -> t (t m) ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    v) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
    forall a b. (a -> b) -> a -> b
    $
            QualName VName -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
    forall {k} {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
           {m :: * -> *} {as}.
    (Ord k, MonadWriter [TypeParam] (t (t m)), MonadFreshNames m,
     MonadTrans t, MonadTrans t,
     MonadState (Map k (RetTypeBase Exp as), Map Int Exp) (t (t m)),
     Monad (t m)) =>
    QualName k -> RetTypeBase MonoSize as -> t (t m) ()
    addSubst QualName VName
    v RetTypeBase MonoSize NoUniqueness
    rt
        subRet StructType
    t1 (RetType [VName]
    _ MonoType
    t2) =
          StructType -> MonoType -> t (t m) ()
    sub StructType
    t1 MonoType
    t2

        sub :: StructType -> MonoType -> t (t m) ()
    sub t1 :: StructType
    t1@(Array NoUniqueness
    _ (Shape (Exp
    d1 : [Exp]
    _)) ScalarTypeBase Exp NoUniqueness
    _) t2 :: MonoType
    t2@(Array NoUniqueness
    _ (Shape (MonoSize
    d2 : [MonoSize]
    _)) ScalarTypeBase MonoSize NoUniqueness
    _) = do
          case MonoSize
    d2 of
            MonoAnon Int
    i -> do
              (Map VName StructRetType
    ts, Map Int Exp
    sizes) <- t (t m) (Map VName StructRetType, Map Int Exp)
    forall s (m :: * -> *). MonadState s m => m s
    get
              (Map VName StructRetType, Map Int Exp) -> t (t m) ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put (Map VName StructRetType
    ts, Int -> Exp -> Map Int Exp -> Map Int Exp
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert Int
    i Exp
    d1 Map Int Exp
    sizes)
            MonoSize
    _ -> () -> t (t m) ()
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
          StructType -> MonoType -> t (t m) ()
    sub (Int -> StructType -> StructType
    forall dim as. Int -> TypeBase dim as -> TypeBase dim as
    stripArray Int
    1 StructType
    t1) (Int -> MonoType -> MonoType
    forall dim as. Int -> TypeBase dim as -> TypeBase dim as
    stripArray Int
    1 MonoType
    t2)
        sub (Scalar (TypeVar NoUniqueness
    _ QualName VName
    v [TypeArg Exp]
    _)) MonoType
    t =
          Bool -> t (t m) () -> t (t m) ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (VName -> Int
    baseTag (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    v) Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    maxIntrinsicTag) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
    forall a b. (a -> b) -> a -> b
    $
            QualName VName -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
    forall {k} {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
           {m :: * -> *} {as}.
    (Ord k, MonadWriter [TypeParam] (t (t m)), MonadFreshNames m,
     MonadTrans t, MonadTrans t,
     MonadState (Map k (RetTypeBase Exp as), Map Int Exp) (t (t m)),
     Monad (t m)) =>
    QualName k -> RetTypeBase MonoSize as -> t (t m) ()
    addSubst QualName VName
    v (RetTypeBase MonoSize NoUniqueness -> t (t m) ())
    -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
    forall a b. (a -> b) -> a -> b
    $
              [VName] -> MonoType -> RetTypeBase MonoSize NoUniqueness
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [] MonoType
    t
        sub (Scalar (Record Map Name StructType
    fields1)) (Scalar (Record Map Name MonoType
    fields2)) =
          (StructType -> MonoType -> t (t m) ())
    -> [StructType] -> [MonoType] -> t (t m) ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_
            StructType -> MonoType -> t (t m) ()
    sub
            (((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
    $ Map Name StructType -> [(Name, StructType)]
    forall a. Map Name a -> [(Name, a)]
    sortFields Map Name StructType
    fields1)
            (((Name, MonoType) -> MonoType) -> [(Name, MonoType)] -> [MonoType]
    forall a b. (a -> b) -> [a] -> [b]
    map (Name, MonoType) -> MonoType
    forall a b. (a, b) -> b
    snd ([(Name, MonoType)] -> [MonoType])
    -> [(Name, MonoType)] -> [MonoType]
    forall a b. (a -> b) -> a -> b
    $ Map Name MonoType -> [(Name, MonoType)]
    forall a. Map Name a -> [(Name, a)]
    sortFields Map Name MonoType
    fields2)
        sub (Scalar Prim {}) (Scalar Prim {}) = () -> t (t m) ()
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        sub (Scalar (Arrow NoUniqueness
    _ PName
    _ Diet
    _ StructType
    t1a (RetType [VName]
    _ TypeBase Exp Uniqueness
    t1b))) (Scalar (Arrow NoUniqueness
    _ PName
    _ Diet
    _ MonoType
    t2a RetTypeBase MonoSize Uniqueness
    t2b)) = do
          StructType -> MonoType -> t (t m) ()
    sub StructType
    t1a MonoType
    t2a
          StructType -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
    subRet (TypeBase Exp Uniqueness -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct TypeBase Exp Uniqueness
    t1b) ((Uniqueness -> NoUniqueness)
    -> RetTypeBase MonoSize Uniqueness
    -> RetTypeBase MonoSize NoUniqueness
    forall b c a. (b -> c) -> RetTypeBase a b -> RetTypeBase a c
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (NoUniqueness -> Uniqueness -> NoUniqueness
    forall a b. a -> b -> a
    const NoUniqueness
    NoUniqueness) RetTypeBase MonoSize Uniqueness
    t2b)
        sub (Scalar (Sum Map Name [StructType]
    cs1)) (Scalar (Sum Map Name [MonoType]
    cs2)) =
          ((Name, [StructType]) -> (Name, [MonoType]) -> t (t m) [()])
    -> [(Name, [StructType])] -> [(Name, [MonoType])] -> t (t m) ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ (Name, [StructType]) -> (Name, [MonoType]) -> t (t m) [()]
    forall {a} {a}.
    (a, [StructType]) -> (a, [MonoType]) -> t (t m) [()]
    typeSubstClause (Map Name [StructType] -> [(Name, [StructType])]
    forall a. Map Name a -> [(Name, a)]
    sortConstrs Map Name [StructType]
    cs1) (Map Name [MonoType] -> [(Name, [MonoType])]
    forall a. Map Name a -> [(Name, a)]
    sortConstrs Map Name [MonoType]
    cs2)
          where
            typeSubstClause :: (a, [StructType]) -> (a, [MonoType]) -> t (t m) [()]
    typeSubstClause (a
    _, [StructType]
    ts1) (a
    _, [MonoType]
    ts2) = (StructType -> MonoType -> t (t m) ())
    -> [StructType] -> [MonoType] -> t (t m) [()]
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m [c]
    zipWithM StructType -> MonoType -> t (t m) ()
    sub [StructType]
    ts1 [MonoType]
    ts2
        sub t1 :: StructType
    t1@(Scalar Sum {}) MonoType
    t2 = StructType -> MonoType -> t (t m) ()
    sub StructType
    t1 MonoType
    t2
        sub StructType
    t1 t2 :: MonoType
    t2@(Scalar Sum {}) = StructType -> MonoType -> t (t m) ()
    sub StructType
    t1 MonoType
    t2
        sub StructType
    t1 MonoType
    t2 = [Char] -> t (t m) ()
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> t (t m) ()) -> [Char] -> t (t m) ()
    forall a b. (a -> b) -> a -> b
    $ [[Char]] -> [Char]
    unlines [[Char]
    "typeSubstsM: mismatched types:", StructType -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString StructType
    t1, MonoType -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString MonoType
    t2]
        addSubst :: QualName k -> RetTypeBase MonoSize as -> t (t m) ()
    addSubst (QualName [k]
    _ k
    v) (RetType [VName]
    ext TypeBase MonoSize as
    t) = do
          (Map k (RetTypeBase Exp as)
    ts, Map Int Exp
    sizes) <- t (t m) (Map k (RetTypeBase Exp as), Map Int Exp)
    forall s (m :: * -> *). MonadState s m => m s
    get
          Bool -> t (t m) () -> t (t m) ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (k
    v k -> Map k (RetTypeBase Exp as) -> Bool
    forall k a. Ord k => k -> Map k a -> Bool
    `M.member` Map k (RetTypeBase Exp as)
    ts) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
    forall a b. (a -> b) -> a -> b
    $ do
            TypeBase Exp as
    t' <- (MonoSize -> t (t m) Exp)
    -> (as -> t (t m) as)
    -> TypeBase MonoSize as
    -> t (t m) (TypeBase Exp as)
    forall (f :: * -> *) a c b d.
    Applicative f =>
    (a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
    forall (t :: * -> * -> *) (f :: * -> *) a c b d.
    (Bitraversable t, Applicative f) =>
    (a -> f c) -> (b -> f d) -> t a b -> f (t c d)
    bitraverse MonoSize -> t (t m) Exp
    forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
           {m :: * -> *} {a}.
    (MonadWriter [TypeParam] (t (t m)), MonadFreshNames m, Monad (t m),
     MonadTrans t, MonadTrans t,
     MonadState (a, Map Int Exp) (t (t m))) =>
    MonoSize -> t (t m) Exp
    onDim as -> t (t m) as
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure TypeBase MonoSize as
    t
            (Map k (RetTypeBase Exp as), Map Int Exp) -> t (t m) ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put (k
    -> RetTypeBase Exp as
    -> Map k (RetTypeBase Exp as)
    -> Map k (RetTypeBase Exp as)
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert k
    v ([VName] -> TypeBase Exp as -> RetTypeBase Exp as
    forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
    RetType [VName]
    ext TypeBase Exp as
    t') Map k (RetTypeBase Exp as)
    ts, Map Int Exp
    sizes)

        onDim :: MonoSize -> t (t m) Exp
    onDim (MonoKnown Int
    i) = do
          (a
    ts, Map Int Exp
    sizes) <- t (t m) (a, Map Int Exp)
    forall s (m :: * -> *). MonadState s m => m s
    get
          case Int -> Map Int Exp -> Maybe Exp
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Int
    i Map Int Exp
    sizes of
            Maybe Exp
    Nothing -> do
              VName
    d <- t m VName -> t (t m) VName
    forall (m :: * -> *) a. Monad m => m a -> t m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (t m VName -> t (t m) VName) -> t m VName -> t (t m) VName
    forall a b. (a -> b) -> a -> b
    $ m VName -> t m VName
    forall (m :: * -> *) a. Monad m => m a -> t m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m VName -> t m VName) -> m VName -> t m VName
    forall a b. (a -> b) -> a -> b
    $ [Char] -> m VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "d"
              [TypeParam] -> t (t m) ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell [VName -> SrcLoc -> TypeParam
    forall vn. vn -> SrcLoc -> TypeParamBase vn
    TypeParamDim VName
    d SrcLoc
    loc]
              (a, Map Int Exp) -> t (t m) ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put (a
    ts, Int -> Exp -> Map Int Exp -> Map Int Exp
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert Int
    i (QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    d) SrcLoc
    forall a. Monoid a => a
    mempty) Map Int Exp
    sizes)
              Exp -> t (t m) Exp
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> t (t m) Exp) -> Exp -> t (t m) Exp
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> SrcLoc -> Exp
    sizeFromName (VName -> QualName VName
    forall v. v -> QualName v
    qualName VName
    d) SrcLoc
    forall a. Monoid a => a
    mempty
            Just Exp
    d ->
              Exp -> t (t m) Exp
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    d
        onDim (MonoAnon Int
    i) = do
          (a
    _, Map Int Exp
    sizes) <- t (t m) (a, Map Int Exp)
    forall s (m :: * -> *). MonadState s m => m s
    get
          case Int -> Map Int Exp -> Maybe Exp
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Int
    i Map Int Exp
    sizes of
            Maybe Exp
    Nothing -> Exp -> t (t m) Exp
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    anySize
            Just Exp
    d -> Exp -> t (t m) Exp
    forall a. a -> t (t m) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    d

    -- Perform a given substitution on the types in a pattern.
    substPat :: (t -> t) -> Pat t -> Pat t
    substPat :: forall t. (t -> t) -> Pat t -> Pat t
    substPat t -> t
    f Pat t
    pat = case Pat t
    pat of
      TuplePat [Pat t]
    pats SrcLoc
    loc -> [Pat t] -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t.
    [PatBase f vn t] -> SrcLoc -> PatBase f vn t
    TuplePat ((Pat t -> Pat t) -> [Pat t] -> [Pat t]
    forall a b. (a -> b) -> [a] -> [b]
    map ((t -> t) -> Pat t -> Pat t
    forall t. (t -> t) -> Pat t -> Pat t
    substPat t -> t
    f) [Pat t]
    pats) SrcLoc
    loc
      RecordPat [(L Name, Pat t)]
    fs SrcLoc
    loc -> [(L Name, Pat t)] -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t.
    [(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
    RecordPat (((L Name, Pat t) -> (L Name, Pat t))
    -> [(L Name, Pat t)] -> [(L Name, Pat t)]
    forall a b. (a -> b) -> [a] -> [b]
    map (L Name, Pat t) -> (L Name, Pat t)
    forall {a}. (a, Pat t) -> (a, Pat t)
    substField [(L Name, Pat t)]
    fs) SrcLoc
    loc
        where
          substField :: (a, Pat t) -> (a, Pat t)
    substField (a
    n, Pat t
    p) = (a
    n, (t -> t) -> Pat t -> Pat t
    forall t. (t -> t) -> Pat t -> Pat t
    substPat t -> t
    f Pat t
    p)
      PatParens Pat t
    p SrcLoc
    loc -> Pat t -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t.
    PatBase f vn t -> SrcLoc -> PatBase f vn t
    PatParens ((t -> t) -> Pat t -> Pat t
    forall t. (t -> t) -> Pat t -> Pat t
    substPat t -> t
    f Pat t
    p) SrcLoc
    loc
      PatAttr AttrInfo VName
    attr Pat t
    p SrcLoc
    loc -> AttrInfo VName -> Pat t -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t.
    AttrInfo vn -> PatBase f vn t -> SrcLoc -> PatBase f vn t
    PatAttr AttrInfo VName
    attr ((t -> t) -> Pat t -> Pat t
    forall t. (t -> t) -> Pat t -> Pat t
    substPat t -> t
    f Pat t
    p) SrcLoc
    loc
      Id VName
    vn (Info t
    tp) SrcLoc
    loc -> VName -> Info t -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
    Id VName
    vn (t -> Info t
    forall a. a -> Info a
    Info (t -> Info t) -> t -> Info t
    forall a b. (a -> b) -> a -> b
    $ t -> t
    f t
    tp) SrcLoc
    loc
      Wildcard (Info t
    tp) SrcLoc
    loc -> Info t -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
    Wildcard (t -> Info t
    forall a. a -> Info a
    Info (t -> Info t) -> t -> Info t
    forall a b. (a -> b) -> a -> b
    $ t -> t
    f t
    tp) SrcLoc
    loc
      PatAscription Pat t
    p TypeExp Exp VName
    _ SrcLoc
    _ -> (t -> t) -> Pat t -> Pat t
    forall t. (t -> t) -> Pat t -> Pat t
    substPat t -> t
    f Pat t
    p
      PatLit PatLit
    e (Info t
    tp) SrcLoc
    loc -> PatLit -> Info t -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t.
    PatLit -> f t -> SrcLoc -> PatBase f vn t
    PatLit PatLit
    e (t -> Info t
    forall a. a -> Info a
    Info (t -> Info t) -> t -> Info t
    forall a b. (a -> b) -> a -> b
    $ t -> t
    f t
    tp) SrcLoc
    loc
      PatConstr Name
    n (Info t
    tp) [Pat t]
    ps SrcLoc
    loc -> Name -> Info t -> [Pat t] -> SrcLoc -> Pat t
    forall (f :: * -> *) vn t.
    Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
    PatConstr Name
    n (t -> Info t
    forall a. a -> Info a
    Info (t -> Info t) -> t -> Info t
    forall a b. (a -> b) -> a -> b
    $ t -> t
    f t
    tp) [Pat t]
    ps SrcLoc
    loc

    toPolyBinding :: ValBind -> PolyBinding
    toPolyBinding :: ValBind -> PolyBinding
    toPolyBinding (ValBind Maybe (Info EntryPoint)
    entry VName
    name Maybe (TypeExp Exp VName)
    _ (Info RetTypeBase Exp Uniqueness
    rettype) [TypeParam]
    tparams [Pat ParamType]
    params Exp
    body Maybe DocComment
    _ [AttrInfo VName]
    attrs SrcLoc
    loc) =
      (Maybe EntryPoint, VName, [TypeParam], [Pat ParamType],
     RetTypeBase Exp Uniqueness, Exp, [AttrInfo VName], SrcLoc)
    -> PolyBinding
    PolyBinding (Info EntryPoint -> EntryPoint
    forall a. Info a -> a
    unInfo (Info EntryPoint -> EntryPoint)
    -> Maybe (Info EntryPoint) -> Maybe EntryPoint
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Maybe (Info EntryPoint)
    entry, VName
    name, [TypeParam]
    tparams, [Pat ParamType]
    params, RetTypeBase Exp Uniqueness
    rettype, Exp
    body, [AttrInfo VName]
    attrs, SrcLoc
    loc)

    transformValBind :: ValBind -> MonoM Env
    transformValBind :: ValBind -> MonoM Env
    transformValBind ValBind
    valbind = do
      let valbind' :: PolyBinding
    valbind' = ValBind -> PolyBinding
    toPolyBinding ValBind
    valbind

      Bool -> MonoM () -> MonoM ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Maybe (Info EntryPoint) -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe (Info EntryPoint) -> Bool)
    -> Maybe (Info EntryPoint) -> Bool
    forall a b. (a -> b) -> a -> b
    $ ValBind -> Maybe (Info EntryPoint)
    forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
    valBindEntryPoint ValBind
    valbind) (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
    forall a b. (a -> b) -> a -> b
    $ do
        let t :: StructType
    t =
              [Pat ParamType] -> RetTypeBase Exp Uniqueness -> StructType
    funType (ValBind -> [Pat ParamType]
    forall (f :: * -> *) vn.
    ValBindBase f vn -> [PatBase f vn ParamType]
    valBindParams ValBind
    valbind) (RetTypeBase Exp Uniqueness -> StructType)
    -> RetTypeBase Exp Uniqueness -> StructType
    forall a b. (a -> b) -> a -> b
    $
                Info (RetTypeBase Exp Uniqueness) -> RetTypeBase Exp Uniqueness
    forall a. Info a -> a
    unInfo (Info (RetTypeBase Exp Uniqueness) -> RetTypeBase Exp Uniqueness)
    -> Info (RetTypeBase Exp Uniqueness) -> RetTypeBase Exp Uniqueness
    forall a b. (a -> b) -> a -> b
    $
                  ValBind -> Info (RetTypeBase Exp Uniqueness)
    forall (f :: * -> *) vn.
    ValBindBase f vn -> f (RetTypeBase Exp Uniqueness)
    valBindRetType ValBind
    valbind
        (VName
    name, InferSizeArgs
    infer, ValBind
    valbind'') <- PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
    monomorphiseBinding PolyBinding
    valbind' (MonoType -> MonoM (VName, InferSizeArgs, ValBind))
    -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
    forall a b. (a -> b) -> a -> b
    $ StructType -> MonoType
    forall als. TypeBase Exp als -> MonoType
    monoType StructType
    t
        Seq (VName, ValBind) -> MonoM ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell (Seq (VName, ValBind) -> MonoM ())
    -> Seq (VName, ValBind) -> MonoM ()
    forall a b. (a -> b) -> a -> b
    $ (VName, ValBind) -> Seq (VName, ValBind)
    forall a. a -> Seq a
    Seq.singleton (VName
    name, ValBind
    valbind'')
        VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
    addLifted (ValBind -> VName
    forall (f :: * -> *) vn. ValBindBase f vn -> vn
    valBindName ValBind
    valbind) (StructType -> MonoType
    forall als. TypeBase Exp als -> MonoType
    monoType StructType
    t) (VName
    name, InferSizeArgs
    infer)

      Env -> MonoM Env
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
        Env
    forall a. Monoid a => a
    mempty
          { envPolyBindings = M.singleton (valBindName valbind) valbind',
            envGlobalScope =
              if null (valBindParams valbind)
                then S.fromList $ retDims $ unInfo $ valBindRetType valbind
                else mempty
          }

    transformValBinds :: [ValBind] -> MonoM ()
    transformValBinds :: [ValBind] -> MonoM ()
    transformValBinds [] = () -> MonoM ()
    forall a. a -> MonoM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    transformValBinds (ValBind
    valbind : [ValBind]
    ds) = do
      Env
    env <- ValBind -> MonoM Env
    transformValBind ValBind
    valbind
      Env -> MonoM () -> MonoM ()
    forall a. Env -> MonoM a -> MonoM a
    localEnv Env
    env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
    forall a b. (a -> b) -> a -> b
    $ [ValBind] -> MonoM ()
    transformValBinds [ValBind]
    ds

    -- | Monomorphise a list of top-level value bindings.
    transformProg :: (MonadFreshNames m) => [ValBind] -> m [ValBind]
    transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
    transformProg [ValBind]
    decs =
      (((), Seq (VName, ValBind)) -> [ValBind])
    -> m ((), Seq (VName, ValBind)) -> m [ValBind]
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Seq ValBind -> [ValBind]
    forall a. Seq a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList (Seq ValBind -> [ValBind])
    -> (((), Seq (VName, ValBind)) -> Seq ValBind)
    -> ((), Seq (VName, ValBind))
    -> [ValBind]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, ValBind) -> ValBind)
    -> Seq (VName, ValBind) -> Seq ValBind
    forall a b. (a -> b) -> Seq a -> Seq b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (VName, ValBind) -> ValBind
    forall a b. (a, b) -> b
    snd (Seq (VName, ValBind) -> Seq ValBind)
    -> (((), Seq (VName, ValBind)) -> Seq (VName, ValBind))
    -> ((), Seq (VName, ValBind))
    -> Seq ValBind
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((), Seq (VName, ValBind)) -> Seq (VName, ValBind)
    forall a b. (a, b) -> b
    snd) (m ((), Seq (VName, ValBind)) -> m [ValBind])
    -> m ((), Seq (VName, ValBind)) -> m [ValBind]
    forall a b. (a -> b) -> a -> b
    $
        (VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
    -> m ((), Seq (VName, ValBind))
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
     -> m ((), Seq (VName, ValBind)))
    -> (VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
    -> m ((), Seq (VName, ValBind))
    forall a b. (a -> b) -> a -> b
    $ \VNameSource
    namesrc ->
          VNameSource
    -> MonoM () -> (((), Seq (VName, ValBind)), VNameSource)
    forall a.
    VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
    runMonoM VNameSource
    namesrc (MonoM () -> (((), Seq (VName, ValBind)), VNameSource))
    -> MonoM () -> (((), Seq (VName, ValBind)), VNameSource)
    forall a b. (a -> b) -> a -> b
    $ [ValBind] -> MonoM ()
    transformValBinds [ValBind]
    decs