    -- | Type inference of @loop@.  This is complicated because of the
    -- uniqueness and size inference, so the implementation is separate
    -- from the main type checker.
    module Language.Futhark.TypeChecker.Terms.Loop
      ( UncheckedLoop,
        CheckedLoop,
        checkLoop,
      )
    where

    import Control.Monad
    import Control.Monad.Reader
    import Control.Monad.State
    import Data.Bifunctor
    import Data.Bitraversable
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Data.Set qualified as S
    import Futhark.Util (nubOrd)
    import Futhark.Util.Pretty hiding (group, space)
    import Language.Futhark
    import Language.Futhark.TypeChecker.Monad hiding (BoundV)
    import Language.Futhark.TypeChecker.Terms.Monad
    import Language.Futhark.TypeChecker.Terms.Pat
    import Language.Futhark.TypeChecker.Types
    import Language.Futhark.TypeChecker.Unify
    import Prelude hiding (mod)

    -- | Retrieve an oracle that can be used to decide whether two are in
    -- the same equivalence class (i.e. have been unified).  This is an
    -- exotic operation.
    getAreSame :: (MonadUnify m) => m (VName -> VName -> Bool)
    getAreSame :: forall (m :: * -> *). MonadUnify m => m (VName -> VName -> Bool)
    getAreSame = Map VName (Level, Constraint) -> VName -> VName -> Bool
    forall {a}. Map VName (a, Constraint) -> VName -> VName -> Bool
    check (Map VName (Level, Constraint) -> VName -> VName -> Bool)
    -> m (Map VName (Level, Constraint)) -> m (VName -> VName -> Bool)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> m (Map VName (Level, Constraint))
    forall (m :: * -> *).
    MonadUnify m =>
    m (Map VName (Level, Constraint))
    getConstraints
      where
        check :: Map VName (a, Constraint) -> VName -> VName -> Bool
    check Map VName (a, Constraint)
    constraints VName
    x VName
    y =
          case (VName -> Map VName (a, Constraint) -> Maybe (a, Constraint)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x Map VName (a, Constraint)
    constraints, VName -> Map VName (a, Constraint) -> Maybe (a, Constraint)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    y Map VName (a, Constraint)
    constraints) of
            (Just (a
    _, Size (Just (Var QualName VName
    x' Info StructType
    _ SrcLoc
    _)) Usage
    _), Maybe (a, Constraint)
    _) ->
              Map VName (a, Constraint) -> VName -> VName -> Bool
    check Map VName (a, Constraint)
    constraints (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    x') VName
    y
            (Maybe (a, Constraint)
    _, Just (a
    _, Size (Just (Var QualName VName
    y' Info StructType
    _ SrcLoc
    _)) Usage
    _)) ->
              Map VName (a, Constraint) -> VName -> VName -> Bool
    check Map VName (a, Constraint)
    constraints VName
    x (QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    y')
            (Maybe (a, Constraint), Maybe (a, Constraint))
    _ ->
              VName
    x VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    y

    -- | Replace specified sizes with distinct fresh size variables.
    someDimsFreshInType ::
      SrcLoc ->
      Name ->
      [VName] ->
      TypeBase Size als ->
      TermTypeM (TypeBase Size als)
    someDimsFreshInType :: forall als.
    SrcLoc
    -> Name
    -> [VName]
    -> TypeBase Exp als
    -> TermTypeM (TypeBase Exp als)
    someDimsFreshInType SrcLoc
    loc Name
    desc [VName]
    fresh TypeBase Exp als
    t = do
      VName -> VName -> Bool
    areSameSize <- TermTypeM (VName -> VName -> Bool)
    forall (m :: * -> *). MonadUnify m => m (VName -> VName -> Bool)
    getAreSame
      let freshen :: VName -> Bool
    freshen VName
    v = (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> VName -> Bool
    areSameSize VName
    v) [VName]
    fresh
      (Exp -> TermTypeM Exp)
    -> (als -> TermTypeM als)
    -> TypeBase Exp als
    -> TermTypeM (TypeBase Exp als)
    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 ((VName -> Bool) -> Exp -> TermTypeM Exp
    forall {m :: * -> *}.
    MonadUnify m =>
    (VName -> Bool) -> Exp -> m Exp
    onDim VName -> Bool
    freshen) als -> TermTypeM als
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure TypeBase Exp als
    t
      where
        onDim :: (VName -> Bool) -> Exp -> m Exp
    onDim VName -> Bool
    freshen (Var QualName VName
    d Info StructType
    _ SrcLoc
    _)
          | VName -> Bool
    freshen (VName -> Bool) -> VName -> Bool
    forall a b. (a -> b) -> a -> b
    $ QualName VName -> VName
    forall vn. QualName vn -> vn
    qualLeaf QualName VName
    d = do
              VName
    v <- Usage -> Name -> m VName
    forall (m :: * -> *). MonadUnify m => Usage -> Name -> m VName
    newFlexibleDim (SrcLoc -> Usage
    forall a. Located a => a -> Usage
    mkUsage' SrcLoc
    loc) Name
    desc
              Exp -> m Exp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp -> m Exp) -> Exp -> 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
    loc
        onDim VName -> Bool
    _ Exp
    d = Exp -> m Exp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    d

    -- | Replace the specified sizes with fresh size variables of the
    -- specified ridigity.  Returns the new fresh size variables.
    freshDimsInType ::
      Usage ->
      Rigidity ->
      Name ->
      [VName] ->
      TypeBase Size u ->
      TermTypeM (TypeBase Size u, [VName])
    freshDimsInType :: forall u.
    Usage
    -> Rigidity
    -> Name
    -> [VName]
    -> TypeBase Exp u
    -> TermTypeM (TypeBase Exp u, [VName])
    freshDimsInType Usage
    usage Rigidity
    r Name
    desc [VName]
    fresh TypeBase Exp u
    t = do
      VName -> VName -> Bool
    areSameSize <- TermTypeM (VName -> VName -> Bool)
    forall (m :: * -> *). MonadUnify m => m (VName -> VName -> Bool)
    getAreSame
      ([(VName, VName)] -> [VName])
    -> (TypeBase Exp u, [(VName, VName)]) -> (TypeBase Exp u, [VName])
    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 (((VName, VName) -> VName) -> [(VName, VName)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, VName) -> VName
    forall a b. (a, b) -> b
    snd) ((TypeBase Exp u, [(VName, VName)]) -> (TypeBase Exp u, [VName]))
    -> TermTypeM (TypeBase Exp u, [(VName, VName)])
    -> TermTypeM (TypeBase Exp u, [VName])
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StateT [(VName, VName)] TermTypeM (TypeBase Exp u)
    -> [(VName, VName)] -> TermTypeM (TypeBase Exp u, [(VName, VName)])
    forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
    runStateT ((Exp -> StateT [(VName, VName)] TermTypeM Exp)
    -> (u -> StateT [(VName, VName)] TermTypeM u)
    -> TypeBase Exp u
    -> StateT [(VName, VName)] TermTypeM (TypeBase Exp u)
    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 ((VName -> VName -> Bool)
    -> Exp -> StateT [(VName, VName)] TermTypeM Exp
    forall {m :: * -> *} {t :: (* -> *) -> * -> *}.
    (MonadUnify m, MonadTrans t, MonadState [(VName, VName)] (t m)) =>
    (VName -> VName -> Bool) -> Exp -> t m Exp
    onDim VName -> VName -> Bool
    areSameSize) u -> StateT [(VName, VName)] TermTypeM u
    forall a. a -> StateT [(VName, VName)] TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure TypeBase Exp u
    t) [(VName, VName)]
    forall a. Monoid a => a
    mempty
      where
        onDim :: (VName -> VName -> Bool) -> Exp -> t m Exp
    onDim VName -> VName -> Bool
    areSameSize (Var (QualName [VName]
    _ VName
    d) Info StructType
    _ SrcLoc
    _)
          | (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> VName -> Bool
    areSameSize VName
    d) [VName]
    fresh = do
              Maybe (VName, VName)
    prev_subst <- ([(VName, VName)] -> Maybe (VName, VName))
    -> t m (Maybe (VName, VName))
    forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
    gets (([(VName, VName)] -> Maybe (VName, VName))
     -> t m (Maybe (VName, VName)))
    -> ([(VName, VName)] -> Maybe (VName, VName))
    -> t m (Maybe (VName, VName))
    forall a b. (a -> b) -> a -> b
    $ ((VName, VName) -> Bool)
    -> [(VName, VName)] -> Maybe (VName, VName)
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    L.find (VName -> VName -> Bool
    areSameSize VName
    d (VName -> Bool)
    -> ((VName, VName) -> VName) -> (VName, VName) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, VName) -> VName
    forall a b. (a, b) -> a
    fst)
              case Maybe (VName, VName)
    prev_subst of
                Just (VName
    _, VName
    d') -> 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
    d') (SrcLoc -> Exp) -> SrcLoc -> Exp
    forall a b. (a -> b) -> a -> b
    $ Usage -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Usage
    usage
                Maybe (VName, VName)
    Nothing -> 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
    $ Usage -> Rigidity -> Name -> m VName
    forall (m :: * -> *).
    MonadUnify m =>
    Usage -> Rigidity -> Name -> m VName
    newDimVar Usage
    usage Rigidity
    r Name
    desc
                  ([(VName, VName)] -> [(VName, VName)]) -> t m ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((VName
    d, 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 -> Exp) -> SrcLoc -> Exp
    forall a b. (a -> b) -> a -> b
    $ Usage -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Usage
    usage
        onDim VName -> VName -> Bool
    _ Exp
    d = Exp -> t m Exp
    forall a. a -> t m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    d

    data ArgSource = Initial | BodyResult

    wellTypedLoopArg :: ArgSource -> [VName] -> Pat ParamType -> Exp -> TermTypeM ()
    wellTypedLoopArg :: ArgSource -> [VName] -> Pat ParamType -> Exp -> TermTypeM ()
    wellTypedLoopArg ArgSource
    src [VName]
    sparams Pat ParamType
    pat Exp
    arg = do
      (StructType
    merge_t, [VName]
    _) <-
        Usage
    -> Rigidity
    -> Name
    -> [VName]
    -> StructType
    -> TermTypeM (StructType, [VName])
    forall u.
    Usage
    -> Rigidity
    -> Name
    -> [VName]
    -> TypeBase Exp u
    -> TermTypeM (TypeBase Exp u, [VName])
    freshDimsInType (Exp -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage Exp
    arg Text
    desc) Rigidity
    Nonrigid Name
    "loop" [VName]
    sparams (StructType -> TermTypeM (StructType, [VName]))
    -> StructType -> TermTypeM (StructType, [VName])
    forall a b. (a -> b) -> a -> b
    $
          ParamType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct (Pat ParamType -> ParamType
    forall d u. Pat (TypeBase d u) -> TypeBase d u
    patternType Pat ParamType
    pat)
      StructType
    arg_t <- StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct (StructType -> StructType)
    -> TermTypeM StructType -> TermTypeM StructType
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Exp -> TermTypeM StructType
    expTypeFully Exp
    arg
      Checking -> TermTypeM () -> TermTypeM ()
    forall a. Checking -> TermTypeM a -> TermTypeM a
    onFailure (StructType -> StructType -> Checking
    checking StructType
    merge_t StructType
    arg_t) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
    forall a b. (a -> b) -> a -> b
    $
        Usage -> StructType -> StructType -> TermTypeM ()
    forall (m :: * -> *).
    MonadUnify m =>
    Usage -> StructType -> StructType -> m ()
    unify (Exp -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage Exp
    arg Text
    desc) StructType
    merge_t StructType
    arg_t
      where
        (StructType -> StructType -> Checking
    checking, Text
    desc) =
          case ArgSource
    src of
            ArgSource
    Initial -> (StructType -> StructType -> Checking
    CheckingLoopInitial, Text
    "matching initial loop values to pattern")
            ArgSource
    BodyResult -> (StructType -> StructType -> Checking
    CheckingLoopBody, Text
    "matching loop body to pattern")

    -- | An un-checked loop.
    type UncheckedLoop =
      (PatBase NoInfo VName ParamType, LoopInitBase NoInfo VName, LoopFormBase NoInfo VName, ExpBase NoInfo VName)

    -- | A loop that has been type-checked.
    type CheckedLoop =
      ([VName], Pat ParamType, LoopInitBase Info VName, LoopFormBase Info VName, Exp)

    checkForImpossible :: Loc -> S.Set VName -> ParamType -> TermTypeM ()
    checkForImpossible :: Loc -> Set VName -> ParamType -> TermTypeM ()
    checkForImpossible Loc
    loc Set VName
    known_before ParamType
    pat_t = do
      Map VName (Level, Constraint)
    cs <- TermTypeM (Map VName (Level, Constraint))
    forall (m :: * -> *).
    MonadUnify m =>
    m (Map VName (Level, Constraint))
    getConstraints
      let bad :: VName -> Maybe (m a)
    bad VName
    v = do
            Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ VName
    v VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.notMember` Set VName
    known_before
            (Level
    _, UnknownSize Loc
    v_loc RigidSource
    _) <- VName -> Map VName (Level, Constraint) -> Maybe (Level, Constraint)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    v Map VName (Level, Constraint)
    cs
            m a -> Maybe (m a)
    forall a. a -> Maybe a
    Just (m a -> Maybe (m a)) -> (Doc () -> m a) -> Doc () -> Maybe (m a)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SrcLoc -> Notes -> Doc () -> m a
    forall loc a. Located loc => loc -> Notes -> Doc () -> m a
    forall (m :: * -> *) loc a.
    (MonadTypeChecker m, Located loc) =>
    loc -> Notes -> Doc () -> m a
    typeError (Loc -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf Loc
    loc) Notes
    forall a. Monoid a => a
    mempty (Doc () -> Maybe (m a)) -> Doc () -> Maybe (m a)
    forall a b. (a -> b) -> a -> b
    $
              Doc ()
    "Inferred type for loop parameter is"
                Doc () -> Doc () -> Doc ()
    forall a. Doc a -> Doc a -> Doc a
    </> Level -> Doc () -> Doc ()
    forall ann. Level -> Doc ann -> Doc ann
    indent Level
    2 (ParamType -> Doc ()
    forall a ann. Pretty a => a -> Doc ann
    forall ann. ParamType -> Doc ann
    pretty ParamType
    pat_t)
                Doc () -> Doc () -> Doc ()
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ()
    "but"
                Doc () -> Doc () -> Doc ()
    forall a. Doc a -> Doc a -> Doc a
    <+> Doc () -> Doc ()
    forall ann. Doc ann -> Doc ann
    dquotes (VName -> Doc ()
    forall a. VName -> Doc a
    forall v a. IsName v => v -> Doc a
    prettyName VName
    v)
                Doc () -> Doc () -> Doc ()
    forall a. Doc a -> Doc a -> Doc a
    <+> Doc ()
    "is an existential size created inside the loop body at"
                Doc () -> Doc () -> Doc ()
    forall a. Doc a -> Doc a -> Doc a
    <+> String -> Doc ()
    forall ann. String -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty (Loc -> Loc -> String
    forall a b. (Located a, Located b) => a -> b -> String
    locStrRel Loc
    loc Loc
    v_loc)
                Doc () -> Doc () -> Doc ()
    forall a. Semigroup a => a -> a -> a
    <> Doc ()
    "."
      case (VName -> Maybe (TermTypeM ())) -> [VName] -> [TermTypeM ()]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe VName -> Maybe (TermTypeM ())
    forall {m :: * -> *} {a}.
    MonadTypeChecker m =>
    VName -> Maybe (m a)
    bad ([VName] -> [TermTypeM ()]) -> [VName] -> [TermTypeM ()]
    forall a b. (a -> b) -> a -> b
    $ Set VName -> [VName]
    forall a. Set a -> [a]
    S.toList (Set VName -> [VName]) -> Set VName -> [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
    $ ParamType -> FV
    forall u. TypeBase Exp u -> FV
    freeInType ParamType
    pat_t of
        TermTypeM ()
    problem : [TermTypeM ()]
    _ -> TermTypeM ()
    problem
        [] -> () -> TermTypeM ()
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    -- | Type-check a @loop@ expression, passing in a function for
    -- type-checking subexpressions.
    checkLoop ::
      (ExpBase NoInfo VName -> TermTypeM Exp) ->
      UncheckedLoop ->
      SrcLoc ->
      TermTypeM (CheckedLoop, AppRes)
    checkLoop :: (ExpBase NoInfo VName -> TermTypeM Exp)
    -> UncheckedLoop -> SrcLoc -> TermTypeM (CheckedLoop, AppRes)
    checkLoop ExpBase NoInfo VName -> TermTypeM Exp
    checkExp (PatBase NoInfo VName ParamType
    mergepat, LoopInitBase NoInfo VName
    loopinit, LoopFormBase NoInfo VName
    form, ExpBase NoInfo VName
    loopbody) SrcLoc
    loc = do
      Exp
    loopinit' <- ExpBase NoInfo VName -> TermTypeM Exp
    checkExp (ExpBase NoInfo VName -> TermTypeM Exp)
    -> ExpBase NoInfo VName -> TermTypeM Exp
    forall a b. (a -> b) -> a -> b
    $ case LoopInitBase NoInfo VName
    loopinit of
        LoopInitExplicit ExpBase NoInfo VName
    e -> ExpBase NoInfo VName
    e
        LoopInitImplicit NoInfo (ExpBase NoInfo VName)
    _ ->
          -- Should have been filled out in Names
          String -> ExpBase NoInfo VName
    forall a. HasCallStack => String -> a
    error String
    "Unspected LoopInitImplicit"
      Set VName
    known_before <- Map VName (Level, Constraint) -> Set VName
    forall k a. Map k a -> Set k
    M.keysSet (Map VName (Level, Constraint) -> Set VName)
    -> TermTypeM (Map VName (Level, Constraint))
    -> TermTypeM (Set VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> TermTypeM (Map VName (Level, Constraint))
    forall (m :: * -> *).
    MonadUnify m =>
    m (Map VName (Level, Constraint))
    getConstraints
      Usage -> Text -> StructType -> TermTypeM ()
    forall (m :: * -> *).
    MonadUnify m =>
    Usage -> Text -> StructType -> m ()
    zeroOrderType
        (Exp -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage Exp
    loopinit' Text
    "use as loop variable")
        Text
    "type used as loop variable"
        (StructType -> TermTypeM ())
    -> (StructType -> StructType) -> StructType -> TermTypeM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct
        (StructType -> TermTypeM ())
    -> TermTypeM StructType -> TermTypeM ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Exp -> TermTypeM StructType
    expTypeFully Exp
    loopinit'

      -- The handling of dimension sizes is a bit intricate, but very
      -- similar to checking a function, followed by checking a call to
      -- it.  The overall procedure is as follows:
      --
      -- (1) All empty dimensions in the merge pattern are instantiated
      -- with nonrigid size variables.  All explicitly specified
      -- dimensions are preserved.
      --
      -- (2) The body of the loop is type-checked.  The result type is
      -- combined with the merge pattern type to determine which sizes are
      -- variant, and these are turned into size parameters for the merge
      -- pattern.
      --
      -- (3) We now conceptually have a function parameter type and
      -- return type.  We check that it can be called with the body type
      -- as argument.
      --
      -- (4) Similarly to (3), we check that the "function" can be
      -- called with the initial merge values as argument.  The result
      -- of this is the type of the loop as a whole.

      (StructType
    merge_t, Map VName Exp
    new_dims_map) <-
        -- dim handling (1)
        Usage
    -> Rigidity
    -> Name
    -> StructType
    -> TermTypeM (StructType, Map VName Exp)
    forall als.
    Usage
    -> Rigidity
    -> Name
    -> TypeBase Exp als
    -> TermTypeM (TypeBase Exp als, Map VName Exp)
    allDimsFreshInType (SrcLoc -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage SrcLoc
    loc Text
    "loop parameter type inference") Rigidity
    Nonrigid Name
    "loop_d"
          (StructType -> TermTypeM (StructType, Map VName Exp))
    -> TermTypeM StructType -> TermTypeM (StructType, Map VName Exp)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Exp -> TermTypeM StructType
    expTypeFully Exp
    loopinit'
      let new_dims_to_initial_dim :: [(VName, Exp)]
    new_dims_to_initial_dim = Map VName Exp -> [(VName, Exp)]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName Exp
    new_dims_map
          new_dims :: [VName]
    new_dims = ((VName, Exp) -> VName) -> [(VName, Exp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, Exp) -> VName
    forall a b. (a, b) -> a
    fst [(VName, Exp)]
    new_dims_to_initial_dim

      -- dim handling (2)
      let checkLoopReturnSize :: Pat ParamType -> Exp -> TermTypeM ([VName], Pat ParamType)
    checkLoopReturnSize Pat ParamType
    mergepat' Exp
    loopbody' = do
            StructType
    loopbody_t <- Exp -> TermTypeM StructType
    expTypeFully Exp
    loopbody'
            ParamType
    mergepat_t <- ParamType -> TermTypeM ParamType
    forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
    normTypeFully (Pat ParamType -> ParamType
    forall d u. Pat (TypeBase d u) -> TypeBase d u
    patternType Pat ParamType
    mergepat')

            let ok_names :: Set VName
    ok_names = Set VName
    known_before Set VName -> Set VName -> Set VName
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList [VName]
    new_dims
            Loc -> Set VName -> ParamType -> TermTypeM ()
    checkForImpossible (PatBase NoInfo VName ParamType -> Loc
    forall a. Located a => a -> Loc
    locOf PatBase NoInfo VName ParamType
    mergepat) Set VName
    ok_names ParamType
    mergepat_t

            ParamType
    pat_t <- SrcLoc -> Name -> [VName] -> ParamType -> TermTypeM ParamType
    forall als.
    SrcLoc
    -> Name
    -> [VName]
    -> TypeBase Exp als
    -> TermTypeM (TypeBase Exp als)
    someDimsFreshInType SrcLoc
    loc Name
    "loop" [VName]
    new_dims ParamType
    mergepat_t

            -- We are ignoring the dimensions here, because any mismatches
            -- should be turned into fresh size variables.
            Checking -> TermTypeM () -> TermTypeM ()
    forall a. Checking -> TermTypeM a -> TermTypeM a
    onFailure (StructType -> StructType -> Checking
    CheckingLoopBody (ParamType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct ParamType
    pat_t) (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct StructType
    loopbody_t)) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
    forall a b. (a -> b) -> a -> b
    $
              Usage -> StructType -> StructType -> TermTypeM ()
    forall (m :: * -> *).
    MonadUnify m =>
    Usage -> StructType -> StructType -> m ()
    unify
                (ExpBase NoInfo VName -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage ExpBase NoInfo VName
    loopbody Text
    "matching loop body to loop pattern")
                (ParamType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct ParamType
    pat_t)
                (StructType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct StructType
    loopbody_t)

            -- Figure out which of the 'new_dims' dimensions are variant.
            -- This works because we know that each dimension from
            -- new_dims in the pattern is unique and distinct.
            VName -> VName -> Bool
    areSameSize <- TermTypeM (VName -> VName -> Bool)
    forall (m :: * -> *). MonadUnify m => m (VName -> VName -> Bool)
    getAreSame
            let onDims :: p -> Exp -> Exp -> f Exp
    onDims p
    _ Exp
    x Exp
    y
                  | Exp
    x Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    == Exp
    y = Exp -> f Exp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    x
                onDims p
    _ Exp
    e Exp
    d = do
                  Set VName -> (VName -> f ()) -> f ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ (FV -> Set VName
    fvVars (FV -> Set VName) -> FV -> Set VName
    forall a b. (a -> b) -> a -> b
    $ Exp -> FV
    freeInExp Exp
    e) ((VName -> f ()) -> f ()) -> (VName -> f ()) -> f ()
    forall a b. (a -> b) -> a -> b
    $ \VName
    v -> do
                    case ((VName, Exp) -> Bool) -> [(VName, Exp)] -> Maybe (VName, Exp)
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    L.find (VName -> VName -> Bool
    areSameSize VName
    v (VName -> Bool) -> ((VName, Exp) -> VName) -> (VName, Exp) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Exp) -> VName
    forall a b. (a, b) -> a
    fst) [(VName, Exp)]
    new_dims_to_initial_dim of
                      Just (VName
    _, Exp
    e') ->
                        if Exp
    e' Exp -> Exp -> Bool
    forall a. Eq a => a -> a -> Bool
    == Exp
    d
                          then (p (Map VName (Subst t)) [VName]
     -> p (Map VName (Subst t)) [VName])
    -> f ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((p (Map VName (Subst t)) [VName]
      -> p (Map VName (Subst t)) [VName])
     -> f ())
    -> (p (Map VName (Subst t)) [VName]
        -> p (Map VName (Subst t)) [VName])
    -> f ()
    forall a b. (a -> b) -> a -> b
    $ (Map VName (Subst t) -> Map VName (Subst t))
    -> p (Map VName (Subst t)) [VName]
    -> p (Map VName (Subst t)) [VName]
    forall a b c. (a -> b) -> p a c -> p b c
    forall (p :: * -> * -> *) a b c.
    Bifunctor p =>
    (a -> b) -> p a c -> p b c
    first ((Map VName (Subst t) -> Map VName (Subst t))
     -> p (Map VName (Subst t)) [VName]
     -> p (Map VName (Subst t)) [VName])
    -> (Map VName (Subst t) -> Map VName (Subst t))
    -> p (Map VName (Subst t)) [VName]
    -> p (Map VName (Subst t)) [VName]
    forall a b. (a -> b) -> a -> b
    $ VName -> Subst t -> Map VName (Subst t) -> Map VName (Subst t)
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    v (Subst t -> Map VName (Subst t) -> Map VName (Subst t))
    -> Subst t -> Map VName (Subst t) -> Map VName (Subst t)
    forall a b. (a -> b) -> a -> b
    $ Exp -> Subst t
    forall t. Exp -> Subst t
    ExpSubst Exp
    e'
                          else
                            Bool -> f () -> f ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (VName
    v VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.member` Set VName
    known_before) (f () -> f ()) -> f () -> f ()
    forall a b. (a -> b) -> a -> b
    $
                              (p (Map VName (Subst t)) [VName]
     -> p (Map VName (Subst t)) [VName])
    -> f ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify (([VName] -> [VName])
    -> p (Map VName (Subst t)) [VName]
    -> p (Map VName (Subst t)) [VName]
    forall b c a. (b -> c) -> p a b -> p a c
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (VName
    v :))
                      Maybe (VName, Exp)
    _ ->
                        () -> f ()
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
                  Exp -> f Exp
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Exp
    e
            StructType
    loopbody_t' <- StructType -> TermTypeM StructType
    forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
    normTypeFully StructType
    loopbody_t
            StructType
    merge_t' <- StructType -> TermTypeM StructType
    forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
    normTypeFully StructType
    merge_t

            let (Map VName (Subst t)
    init_substs, [VName]
    sparams) =
                  State (Map VName (Subst t), [VName]) StructType
    -> (Map VName (Subst t), [VName]) -> (Map VName (Subst t), [VName])
    forall s a. State s a -> s -> s
    execState (([VName]
     -> Exp
     -> Exp
     -> StateT (Map VName (Subst t), [VName]) Identity Exp)
    -> StructType
    -> StructType
    -> State (Map VName (Subst t), [VName]) StructType
    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 (Map VName (Subst t), [VName]) Identity Exp
    forall {f :: * -> *} {p :: * -> * -> *} {t} {p}.
    (Bifunctor p, MonadState (p (Map VName (Subst t)) [VName]) f) =>
    p -> Exp -> Exp -> f Exp
    onDims StructType
    merge_t' StructType
    loopbody_t') (Map VName (Subst t), [VName])
    forall a. Monoid a => a
    mempty

            -- Make sure that any of new_dims that are invariant will be
            -- replaced with the invariant size in the loop body.  Failure
            -- to do this can cause type annotations to still refer to
            -- new_dims.
            let dimToInit :: (VName, Subst t) -> TermTypeM ()
    dimToInit (VName
    v, ExpSubst Exp
    e) =
                  VName -> Constraint -> TermTypeM ()
    constrain VName
    v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
    forall a b. (a -> b) -> a -> b
    $ Maybe Exp -> Usage -> Constraint
    Size (Exp -> Maybe Exp
    forall a. a -> Maybe a
    Just Exp
    e) (SrcLoc -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage SrcLoc
    loc Text
    "size of loop parameter")
                dimToInit (VName, Subst t)
    _ =
                  () -> TermTypeM ()
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
            ((VName, Subst Any) -> TermTypeM ())
    -> [(VName, Subst Any)] -> TermTypeM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (VName, Subst Any) -> TermTypeM ()
    forall {t}. (VName, Subst t) -> TermTypeM ()
    dimToInit ([(VName, Subst Any)] -> TermTypeM ())
    -> [(VName, Subst Any)] -> TermTypeM ()
    forall a b. (a -> b) -> a -> b
    $ Map VName (Subst Any) -> [(VName, Subst Any)]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName (Subst Any)
    forall {t}. Map VName (Subst t)
    init_substs

            Pat ParamType
    mergepat'' <- TypeSubs -> Pat ParamType -> Pat ParamType
    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)
    forall {t}. Map VName (Subst t)
    init_substs) (Pat ParamType -> Pat ParamType)
    -> TermTypeM (Pat ParamType) -> TermTypeM (Pat ParamType)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Pat ParamType -> TermTypeM (Pat ParamType)
    forall e. ASTMappable e => e -> TermTypeM e
    updateTypes Pat ParamType
    mergepat'

            -- Eliminate those new_dims that turned into sparams so it won't
            -- look like we have ambiguous sizes lying around.
            (Map VName (Level, Constraint) -> Map VName (Level, Constraint))
    -> TermTypeM ()
    forall (m :: * -> *).
    MonadUnify m =>
    (Map VName (Level, Constraint) -> Map VName (Level, Constraint))
    -> m ()
    modifyConstraints ((Map VName (Level, Constraint) -> Map VName (Level, Constraint))
     -> TermTypeM ())
    -> (Map VName (Level, Constraint) -> Map VName (Level, Constraint))
    -> TermTypeM ()
    forall a b. (a -> b) -> a -> b
    $ (VName -> (Level, Constraint) -> Bool)
    -> Map VName (Level, Constraint) -> Map VName (Level, Constraint)
    forall k a. (k -> a -> Bool) -> Map k a -> Map k a
    M.filterWithKey ((VName -> (Level, Constraint) -> Bool)
     -> Map VName (Level, Constraint) -> Map VName (Level, Constraint))
    -> (VName -> (Level, Constraint) -> Bool)
    -> Map VName (Level, Constraint)
    -> Map VName (Level, Constraint)
    forall a b. (a -> b) -> a -> b
    $ \VName
    k (Level, Constraint)
    _ -> VName
    k VName -> [VName] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` [VName]
    sparams

            -- dim handling (3)
            --
            -- The only trick here is that we have to turn any instances
            -- of loop parameters in the type of loopbody' rigid,
            -- because we are no longer in a position to change them,
            -- really.
            ArgSource -> [VName] -> Pat ParamType -> Exp -> TermTypeM ()
    wellTypedLoopArg ArgSource
    BodyResult [VName]
    sparams Pat ParamType
    mergepat'' Exp
    loopbody'

            ([VName], Pat ParamType) -> TermTypeM ([VName], Pat ParamType)
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([VName] -> [VName]
    forall a. Ord a => [a] -> [a]
    nubOrd [VName]
    sparams, Pat ParamType
    mergepat'')

      ([VName]
    sparams, Pat ParamType
    mergepat', LoopFormBase Info VName
    form', Exp
    loopbody') <-
        case LoopFormBase NoInfo VName
    form of
          For IdentBase NoInfo VName StructType
    i ExpBase NoInfo VName
    uboundexp -> do
            Exp
    uboundexp' <-
              Text -> [PrimType] -> Exp -> TermTypeM Exp
    require Text
    "being the bound in a 'for' loop" [PrimType]
    anySignedType
                (Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< ExpBase NoInfo VName -> TermTypeM Exp
    checkExp ExpBase NoInfo VName
    uboundexp
            StructType
    bound_t <- Exp -> TermTypeM StructType
    expTypeFully Exp
    uboundexp'
            IdentBase NoInfo VName StructType
    -> StructType
    -> (Ident StructType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a.
    IdentBase NoInfo VName StructType
    -> StructType -> (Ident StructType -> TermTypeM a) -> TermTypeM a
    bindingIdent IdentBase NoInfo VName StructType
    i StructType
    bound_t ((Ident StructType
      -> TermTypeM
           ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> (Ident StructType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ \Ident StructType
    i' ->
              [SizeBinder VName]
    -> PatBase NoInfo VName ParamType
    -> StructType
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall u a.
    [SizeBinder VName]
    -> PatBase NoInfo VName (TypeBase Exp u)
    -> StructType
    -> (Pat ParamType -> TermTypeM a)
    -> TermTypeM a
    bindingPat [] PatBase NoInfo VName ParamType
    mergepat StructType
    merge_t ((Pat ParamType
      -> TermTypeM
           ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ \Pat ParamType
    mergepat' -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a. TermTypeM a -> TermTypeM a
    incLevel (TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ do
                Exp
    loopbody' <- ExpBase NoInfo VName -> TermTypeM Exp
    checkExp ExpBase NoInfo VName
    loopbody
                ([VName]
    sparams, Pat ParamType
    mergepat'') <- Pat ParamType -> Exp -> TermTypeM ([VName], Pat ParamType)
    checkLoopReturnSize Pat ParamType
    mergepat' Exp
    loopbody'
                ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                  ( [VName]
    sparams,
                    Pat ParamType
    mergepat'',
                    Ident StructType -> Exp -> LoopFormBase Info VName
    forall (f :: * -> *) vn.
    IdentBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
    For Ident StructType
    i' Exp
    uboundexp',
                    Exp
    loopbody'
                  )
          ForIn PatBase NoInfo VName StructType
    xpat ExpBase NoInfo VName
    e -> do
            (StructType
    arr_t, StructType
    _) <- Usage -> Name -> Level -> TermTypeM (StructType, StructType)
    newArrayType (SrcLoc -> Usage
    forall a. Located a => a -> Usage
    mkUsage' (ExpBase NoInfo VName -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf ExpBase NoInfo VName
    e)) Name
    "e" Level
    1
            Exp
    e' <- Text -> StructType -> Exp -> TermTypeM Exp
    unifies Text
    "being iterated in a 'for-in' loop" StructType
    arr_t (Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< ExpBase NoInfo VName -> TermTypeM Exp
    checkExp ExpBase NoInfo VName
    e
            StructType
    t <- Exp -> TermTypeM StructType
    expTypeFully Exp
    e'
            case StructType
    t of
              StructType
    _
                | Just StructType
    t' <- Level -> StructType -> Maybe StructType
    forall dim u. Level -> TypeBase dim u -> Maybe (TypeBase dim u)
    peelArray Level
    1 StructType
    t ->
                    [SizeBinder VName]
    -> PatBase NoInfo VName StructType
    -> StructType
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall u a.
    [SizeBinder VName]
    -> PatBase NoInfo VName (TypeBase Exp u)
    -> StructType
    -> (Pat ParamType -> TermTypeM a)
    -> TermTypeM a
    bindingPat [] PatBase NoInfo VName StructType
    xpat StructType
    t' ((Pat ParamType
      -> TermTypeM
           ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ \Pat ParamType
    xpat' ->
                      [SizeBinder VName]
    -> PatBase NoInfo VName ParamType
    -> StructType
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall u a.
    [SizeBinder VName]
    -> PatBase NoInfo VName (TypeBase Exp u)
    -> StructType
    -> (Pat ParamType -> TermTypeM a)
    -> TermTypeM a
    bindingPat [] PatBase NoInfo VName ParamType
    mergepat StructType
    merge_t ((Pat ParamType
      -> TermTypeM
           ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ \Pat ParamType
    mergepat' -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a. TermTypeM a -> TermTypeM a
    incLevel (TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ do
                        Exp
    loopbody' <- ExpBase NoInfo VName -> TermTypeM Exp
    checkExp ExpBase NoInfo VName
    loopbody
                        ([VName]
    sparams, Pat ParamType
    mergepat'') <- Pat ParamType -> Exp -> TermTypeM ([VName], Pat ParamType)
    checkLoopReturnSize Pat ParamType
    mergepat' Exp
    loopbody'
                        ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                          ( [VName]
    sparams,
                            Pat ParamType
    mergepat'',
                            PatBase Info VName StructType -> Exp -> LoopFormBase Info VName
    forall (f :: * -> *) vn.
    PatBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
    ForIn ((ParamType -> StructType)
    -> Pat ParamType -> 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 ParamType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct Pat ParamType
    xpat') Exp
    e',
                            Exp
    loopbody'
                          )
                | Bool
    otherwise ->
                    SrcLoc
    -> Notes
    -> Doc ()
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM a
    forall (m :: * -> *) loc a.
    (MonadTypeChecker m, Located loc) =>
    loc -> Notes -> Doc () -> m a
    typeError (ExpBase NoInfo VName -> SrcLoc
    forall a. Located a => a -> SrcLoc
    srclocOf ExpBase NoInfo VName
    e) Notes
    forall a. Monoid a => a
    mempty (Doc ()
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> Doc ()
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $
                      Doc ()
    "Iteratee of a for-in loop must be an array, but expression has type"
                        Doc () -> Doc () -> Doc ()
    forall a. Doc a -> Doc a -> Doc a
    <+> StructType -> Doc ()
    forall a ann. Pretty a => a -> Doc ann
    forall ann. StructType -> Doc ann
    pretty StructType
    t
          While ExpBase NoInfo VName
    cond ->
            [SizeBinder VName]
    -> PatBase NoInfo VName ParamType
    -> StructType
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall u a.
    [SizeBinder VName]
    -> PatBase NoInfo VName (TypeBase Exp u)
    -> StructType
    -> (Pat ParamType -> TermTypeM a)
    -> TermTypeM a
    bindingPat [] PatBase NoInfo VName ParamType
    mergepat StructType
    merge_t ((Pat ParamType
      -> TermTypeM
           ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> (Pat ParamType
        -> TermTypeM
             ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ \Pat ParamType
    mergepat' ->
              TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a. TermTypeM a -> TermTypeM a
    incLevel (TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
     -> TermTypeM
          ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a b. (a -> b) -> a -> b
    $ do
                Exp
    cond' <-
                  ExpBase NoInfo VName -> TermTypeM Exp
    checkExp ExpBase NoInfo VName
    cond
                    TermTypeM Exp -> (Exp -> TermTypeM Exp) -> TermTypeM Exp
    forall a b. TermTypeM a -> (a -> TermTypeM b) -> TermTypeM b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Text -> StructType -> Exp -> TermTypeM Exp
    unifies Text
    "being the condition of a 'while' loop" (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
    $ PrimType -> ScalarTypeBase Exp NoUniqueness
    forall dim u. PrimType -> ScalarTypeBase dim u
    Prim PrimType
    Bool)
                Exp
    loopbody' <- ExpBase NoInfo VName -> TermTypeM Exp
    checkExp ExpBase NoInfo VName
    loopbody
                ([VName]
    sparams, Pat ParamType
    mergepat'') <- Pat ParamType -> Exp -> TermTypeM ([VName], Pat ParamType)
    checkLoopReturnSize Pat ParamType
    mergepat' Exp
    loopbody'
                ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    -> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                  ( [VName]
    sparams,
                    Pat ParamType
    mergepat'',
                    Exp -> LoopFormBase Info VName
    forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
    While Exp
    cond',
                    Exp
    loopbody'
                  )

      -- dim handling (4)
      ArgSource -> [VName] -> Pat ParamType -> Exp -> TermTypeM ()
    wellTypedLoopArg ArgSource
    Initial [VName]
    sparams Pat ParamType
    mergepat' Exp
    loopinit'

      (ParamType
    loopt, [VName]
    retext) <-
        Usage
    -> Rigidity
    -> Name
    -> [VName]
    -> ParamType
    -> TermTypeM (ParamType, [VName])
    forall u.
    Usage
    -> Rigidity
    -> Name
    -> [VName]
    -> TypeBase Exp u
    -> TermTypeM (TypeBase Exp u, [VName])
    freshDimsInType
          (SrcLoc -> Text -> Usage
    forall a. Located a => a -> Text -> Usage
    mkUsage SrcLoc
    loc Text
    "inference of loop result type")
          (RigidSource -> Rigidity
    Rigid RigidSource
    RigidLoop)
          Name
    "loop"
          [VName]
    sparams
          (Pat ParamType -> ParamType
    forall d u. Pat (TypeBase d u) -> TypeBase d u
    patternType Pat ParamType
    mergepat')
      (CheckedLoop, AppRes) -> TermTypeM (CheckedLoop, AppRes)
    forall a. a -> TermTypeM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
        ( ([VName]
    sparams, Pat ParamType
    mergepat', Exp -> LoopInitBase Info VName
    forall (f :: * -> *) vn. ExpBase f vn -> LoopInitBase f vn
    LoopInitExplicit Exp
    loopinit', LoopFormBase Info VName
    form', Exp
    loopbody'),
          StructType -> [VName] -> AppRes
    AppRes (ParamType -> StructType
    forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
    toStruct ParamType
    loopt) [VName]
    retext
        )