{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}

-- | 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.DoLoop
  ( UncheckedLoop,
    CheckedLoop,
    checkDoLoop,
  )
where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import Futhark.Util (nubOrd)
import Futhark.Util.Pretty hiding (bool, 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 hiding (Usage)
import Prelude hiding (mod)

-- | Replace specified sizes with distinct fresh size variables.
someDimsFreshInType ::
  SrcLoc ->
  Rigidity ->
  Name ->
  S.Set VName ->
  TypeBase (DimDecl VName) als ->
  TermTypeM (TypeBase (DimDecl VName) als)
someDimsFreshInType :: SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) als
-> TermTypeM (TypeBase (DimDecl VName) als)
someDimsFreshInType SrcLoc
loc Rigidity
r Name
desc Set VName
sizes = (DimDecl VName -> TermTypeM (DimDecl VName))
-> (als -> TermTypeM als)
-> TypeBase (DimDecl VName) als
-> TermTypeM (TypeBase (DimDecl VName) als)
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 DimDecl VName -> TermTypeM (DimDecl VName)
forall (m :: * -> *).
MonadUnify m =>
DimDecl VName -> m (DimDecl VName)
onDim als -> TermTypeM als
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  where
    onDim :: DimDecl VName -> m (DimDecl VName)
onDim (NamedDim QualName VName
d)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
sizes = do
        VName
v <- SrcLoc -> Rigidity -> Name -> m VName
forall (m :: * -> *).
MonadUnify m =>
SrcLoc -> Rigidity -> Name -> m VName
newDimVar SrcLoc
loc Rigidity
r Name
desc
        DimDecl VName -> m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> m (DimDecl VName))
-> DimDecl VName -> m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v
    onDim DimDecl VName
d = DimDecl VName -> m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
d

-- | Replace the specified sizes with fresh size variables of the
-- specified ridigity.  Returns the new fresh size variables.
freshDimsInType ::
  SrcLoc ->
  Rigidity ->
  Name ->
  S.Set VName ->
  TypeBase (DimDecl VName) als ->
  TermTypeM (TypeBase (DimDecl VName) als, [VName])
freshDimsInType :: SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) als
-> TermTypeM (TypeBase (DimDecl VName) als, [VName])
freshDimsInType SrcLoc
loc Rigidity
r Name
desc Set VName
sizes TypeBase (DimDecl VName) als
t =
  (Map VName VName -> [VName])
-> (TypeBase (DimDecl VName) als, Map VName VName)
-> (TypeBase (DimDecl VName) als, [VName])
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Map VName VName -> [VName]
forall k a. Map k a -> [a]
M.elems ((TypeBase (DimDecl VName) als, Map VName VName)
 -> (TypeBase (DimDecl VName) als, [VName]))
-> TermTypeM (TypeBase (DimDecl VName) als, Map VName VName)
-> TermTypeM (TypeBase (DimDecl VName) als, [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (Map VName VName) TermTypeM (TypeBase (DimDecl VName) als)
-> Map VName VName
-> TermTypeM (TypeBase (DimDecl VName) als, Map VName VName)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((DimDecl VName
 -> StateT (Map VName VName) TermTypeM (DimDecl VName))
-> (als -> StateT (Map VName VName) TermTypeM als)
-> TypeBase (DimDecl VName) als
-> StateT
     (Map VName VName) TermTypeM (TypeBase (DimDecl VName) als)
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 DimDecl VName -> StateT (Map VName VName) TermTypeM (DimDecl VName)
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadState (Map VName VName) (t m), MonadTrans t, MonadUnify m) =>
DimDecl VName -> t m (DimDecl VName)
onDim als -> StateT (Map VName VName) TermTypeM als
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase (DimDecl VName) als
t) Map VName VName
forall a. Monoid a => a
mempty
  where
    onDim :: DimDecl VName -> t m (DimDecl VName)
onDim (NamedDim QualName VName
d)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
sizes = do
        Maybe VName
prev_subst <- (Map VName VName -> Maybe VName) -> t m (Maybe VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Map VName VName -> Maybe VName) -> t m (Maybe VName))
-> (Map VName VName -> Maybe VName) -> t m (Maybe VName)
forall a b. (a -> b) -> a -> b
$ VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (VName -> Map VName VName -> Maybe VName)
-> VName -> Map VName VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
        case Maybe VName
prev_subst of
          Just VName
d' -> DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d'
          Maybe VName
Nothing -> do
            VName
v <- m VName -> t m VName
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
$ SrcLoc -> Rigidity -> Name -> m VName
forall (m :: * -> *).
MonadUnify m =>
SrcLoc -> Rigidity -> Name -> m VName
newDimVar SrcLoc
loc Rigidity
r Name
desc
            (Map VName VName -> Map VName VName) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map VName VName -> Map VName VName) -> t m ())
-> (Map VName VName -> Map VName VName) -> t m ()
forall a b. (a -> b) -> a -> b
$ VName -> VName -> Map VName VName -> Map VName VName
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
d) VName
v
            DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v
    onDim DimDecl VName
d = DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
d

-- | An un-checked loop.
type UncheckedLoop =
  (UncheckedPat, UncheckedExp, LoopFormBase NoInfo Name, UncheckedExp)

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

-- | Type-check a @loop@ expression, passing in a function for
-- type-checking subexpressions.
checkDoLoop ::
  (UncheckedExp -> TermTypeM Exp) ->
  UncheckedLoop ->
  SrcLoc ->
  TermTypeM (CheckedLoop, AppRes)
checkDoLoop :: (UncheckedExp -> TermTypeM Exp)
-> UncheckedLoop -> SrcLoc -> TermTypeM (CheckedLoop, AppRes)
checkDoLoop UncheckedExp -> TermTypeM Exp
checkExp (UncheckedPat
mergepat, UncheckedExp
mergeexp, LoopFormBase NoInfo Name
form, UncheckedExp
loopbody) SrcLoc
loc =
  TermTypeM Exp
-> (Exp -> Occurrences -> TermTypeM (CheckedLoop, AppRes))
-> TermTypeM (CheckedLoop, AppRes)
forall a b.
TermTypeM a -> (a -> Occurrences -> TermTypeM b) -> TermTypeM b
sequentially (UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
mergeexp) ((Exp -> Occurrences -> TermTypeM (CheckedLoop, AppRes))
 -> TermTypeM (CheckedLoop, AppRes))
-> (Exp -> Occurrences -> TermTypeM (CheckedLoop, AppRes))
-> TermTypeM (CheckedLoop, AppRes)
forall a b. (a -> b) -> a -> b
$ \Exp
mergeexp' Occurrences
_ -> do
    Usage
-> String -> TypeBase (DimDecl VName) Aliasing -> TermTypeM ()
forall (m :: * -> *) dim as.
(MonadUnify m, Pretty (ShapeDecl dim), Monoid as) =>
Usage -> String -> TypeBase dim as -> m ()
zeroOrderType
      (SrcLoc -> String -> Usage
mkUsage (UncheckedExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf UncheckedExp
mergeexp) String
"use as loop variable")
      String
"type used as loop variable"
      (TypeBase (DimDecl VName) Aliasing -> TermTypeM ())
-> TermTypeM (TypeBase (DimDecl VName) Aliasing) -> TermTypeM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
mergeexp'

    -- 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 initial merge
    -- values as argument.  The result of this is the type of the loop
    -- as a whole.
    --
    -- (There is also a convergence loop for inferring uniqueness, but
    -- that's orthogonal to the size handling.)

    (TypeBase (DimDecl VName) Aliasing
merge_t, Map VName (DimDecl VName)
new_dims_to_initial_dim) <-
      -- dim handling (1)
      SrcLoc
-> Rigidity
-> Name
-> TypeBase (DimDecl VName) Aliasing
-> TermTypeM
     (TypeBase (DimDecl VName) Aliasing, Map VName (DimDecl VName))
forall als.
SrcLoc
-> Rigidity
-> Name
-> TypeBase (DimDecl VName) als
-> TermTypeM
     (TypeBase (DimDecl VName) als, Map VName (DimDecl VName))
allDimsFreshInType SrcLoc
loc Rigidity
Nonrigid Name
"loop" (TypeBase (DimDecl VName) Aliasing
 -> TermTypeM
      (TypeBase (DimDecl VName) Aliasing, Map VName (DimDecl VName)))
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
-> TermTypeM
     (TypeBase (DimDecl VName) Aliasing, Map VName (DimDecl VName))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
mergeexp'
    let new_dims :: [VName]
new_dims = Map VName (DimDecl VName) -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName (DimDecl VName)
new_dims_to_initial_dim

    -- dim handling (2)
    let checkLoopReturnSize :: PatBase Info VName
-> Exp -> TermTypeM ([VName], PatBase Info VName)
checkLoopReturnSize PatBase Info VName
mergepat' Exp
loopbody' = do
          TypeBase (DimDecl VName) Aliasing
loopbody_t <- Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
loopbody'
          TypeBase (DimDecl VName) Aliasing
pat_t <-
            SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) Aliasing
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
forall als.
SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) als
-> TermTypeM (TypeBase (DimDecl VName) als)
someDimsFreshInType SrcLoc
loc Rigidity
Nonrigid Name
"loop" ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
new_dims)
              (TypeBase (DimDecl VName) Aliasing
 -> TermTypeM (TypeBase (DimDecl VName) Aliasing))
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TypeBase (DimDecl VName) Aliasing
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully (PatBase Info VName -> TypeBase (DimDecl VName) Aliasing
patternType PatBase Info VName
mergepat')

          -- 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 (TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
pat_t) (TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
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
              (SrcLoc -> String -> Usage
mkUsage (UncheckedExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf UncheckedExp
loopbody) String
"matching loop body to loop pattern")
              (TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
pat_t)
              (TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
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.
          let onDims :: p -> DimDecl VName -> DimDecl VName -> f (DimDecl VName)
onDims p
_ DimDecl VName
x DimDecl VName
y
                | DimDecl VName
x DimDecl VName -> DimDecl VName -> Bool
forall a. Eq a => a -> a -> Bool
== DimDecl VName
y = DimDecl VName -> f (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
x
              onDims p
_ (NamedDim QualName VName
v) DimDecl VName
d
                | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
new_dims = do
                  case VName -> Map VName (DimDecl VName) -> Maybe (DimDecl VName)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName (DimDecl VName)
new_dims_to_initial_dim of
                    Just DimDecl VName
d'
                      | DimDecl VName
d' DimDecl VName -> DimDecl VName -> Bool
forall a. Eq a => a -> a -> Bool
== DimDecl VName
d ->
                        (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 (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 (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) (DimDecl VName -> Subst t
forall t. DimDecl VName -> Subst t
SizeSubst DimDecl VName
d)
                    Maybe (DimDecl VName)
_ ->
                      (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
$ ([VName] -> [VName])
-> p (Map VName (Subst t)) [VName]
-> p (Map VName (Subst t)) [VName]
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:)
                  DimDecl VName -> f (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> f (DimDecl VName))
-> DimDecl VName -> f (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v
              onDims p
_ DimDecl VName
x DimDecl VName
_ = DimDecl VName -> f (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
x
          TypeBase (DimDecl VName) Aliasing
loopbody_t' <- TypeBase (DimDecl VName) Aliasing
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully TypeBase (DimDecl VName) Aliasing
loopbody_t
          TypeBase (DimDecl VName) Aliasing
merge_t' <- TypeBase (DimDecl VName) Aliasing
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully TypeBase (DimDecl VName) Aliasing
merge_t
          let (Map VName (Subst t)
init_substs, [VName]
sparams) =
                State
  (Map VName (Subst t), [VName]) (TypeBase (DimDecl VName) Aliasing)
-> (Map VName (Subst t), [VName]) -> (Map VName (Subst t), [VName])
forall s a. State s a -> s -> s
execState (([VName]
 -> DimDecl VName
 -> DimDecl VName
 -> StateT (Map VName (Subst t), [VName]) Identity (DimDecl VName))
-> TypeBase (DimDecl VName) Aliasing
-> TypeBase (DimDecl VName) Aliasing
-> State
     (Map VName (Subst t), [VName]) (TypeBase (DimDecl VName) Aliasing)
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]
-> DimDecl VName
-> DimDecl VName
-> StateT (Map VName (Subst t), [VName]) Identity (DimDecl VName)
forall (f :: * -> *) (p :: * -> * -> *) t p.
(Bifunctor p, MonadState (p (Map VName (Subst t)) [VName]) f) =>
p -> DimDecl VName -> DimDecl VName -> f (DimDecl VName)
onDims TypeBase (DimDecl VName) Aliasing
merge_t' TypeBase (DimDecl VName) Aliasing
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, SizeSubst DimDecl VName
d) =
                VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Maybe (DimDecl VName) -> Usage -> Constraint
Size (DimDecl VName -> Maybe (DimDecl VName)
forall a. a -> Maybe a
Just DimDecl VName
d) (SrcLoc -> String -> Usage
mkUsage SrcLoc
loc String
"size of loop parameter")
              dimToInit (VName, Subst t)
_ =
                () -> TermTypeM ()
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

          PatBase Info VName
mergepat'' <- TypeSubs -> PatBase Info VName -> PatBase Info VName
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) (PatBase Info VName -> PatBase Info VName)
-> TermTypeM (PatBase Info VName) -> TermTypeM (PatBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName -> TermTypeM (PatBase Info VName)
forall e. ASTMappable e => e -> TermTypeM e
updateTypes PatBase Info VName
mergepat'

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

    -- First we do a basic check of the loop body to figure out which of
    -- the merge parameters are being consumed.  For this, we first need
    -- to check the merge pattern, which requires the (initial) merge
    -- expression.
    --
    -- Play a little with occurences to ensure it does not look like
    -- none of the merge variables are being used.
    (([VName]
sparams, PatBase Info VName
mergepat', LoopFormBase Info VName
form', Exp
loopbody'), Occurrences
bodyflow) <-
      case LoopFormBase NoInfo Name
form of
        For IdentBase NoInfo Name
i UncheckedExp
uboundexp -> do
          Exp
uboundexp' <-
            String -> [PrimType] -> Exp -> TermTypeM Exp
require String
"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
=<< UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
uboundexp
          TypeBase (DimDecl VName) Aliasing
bound_t <- Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
uboundexp'
          IdentBase NoInfo Name
-> TypeBase (DimDecl VName) Aliasing
-> (Ident
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a.
IdentBase NoInfo Name
-> TypeBase (DimDecl VName) Aliasing
-> (Ident -> TermTypeM a)
-> TermTypeM a
bindingIdent IdentBase NoInfo Name
i TypeBase (DimDecl VName) Aliasing
bound_t ((Ident
  -> TermTypeM
       (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
        Occurrences))
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (Ident
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$ \Ident
i' ->
            TermTypeM
  (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
   Occurrences)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM a
noUnique (TermTypeM
   (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
    Occurrences)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> ((PatBase Info VName
     -> TermTypeM
          (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
           Occurrences))
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a.
[SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName -> TermTypeM a)
-> TermTypeM a
bindingPat [] UncheckedPat
mergepat (TypeBase (DimDecl VName) Aliasing -> InferredType
Ascribed TypeBase (DimDecl VName) Aliasing
merge_t) ((PatBase Info VName
  -> TermTypeM
       (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
        Occurrences))
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$
              \PatBase Info VName
mergepat' -> TermTypeM
  (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
   Occurrences)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM a
onlySelfAliasing (TermTypeM
   (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
    Occurrences)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (TermTypeM
      ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeM
  ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM (a, Occurrences)
tapOccurrences (TermTypeM
   ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$ do
                Exp
loopbody' <- TermTypeM Exp -> TermTypeM Exp
forall a. TermTypeM a -> TermTypeM a
noSizeEscape (TermTypeM Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
loopbody
                ([VName]
sparams, PatBase Info VName
mergepat'') <- PatBase Info VName
-> Exp -> TermTypeM ([VName], PatBase Info VName)
checkLoopReturnSize PatBase Info VName
mergepat' Exp
loopbody'
                ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
forall (m :: * -> *) a. Monad m => a -> m a
return
                  ( [VName]
sparams,
                    PatBase Info VName
mergepat'',
                    Ident -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For Ident
i' Exp
uboundexp',
                    Exp
loopbody'
                  )
        ForIn UncheckedPat
xpat UncheckedExp
e -> do
          (StructType
arr_t, StructType
_) <- SrcLoc -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType (UncheckedExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf UncheckedExp
e) Name
"e" Int
1
          Exp
e' <- String -> StructType -> Exp -> TermTypeM Exp
unifies String
"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
=<< UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
e
          TypeBase (DimDecl VName) Aliasing
t <- Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
e'
          case TypeBase (DimDecl VName) Aliasing
t of
            TypeBase (DimDecl VName) Aliasing
_
              | Just TypeBase (DimDecl VName) Aliasing
t' <- Int
-> TypeBase (DimDecl VName) Aliasing
-> Maybe (TypeBase (DimDecl VName) Aliasing)
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray Int
1 TypeBase (DimDecl VName) Aliasing
t ->
                [SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a.
[SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName -> TermTypeM a)
-> TermTypeM a
bindingPat [] UncheckedPat
xpat (TypeBase (DimDecl VName) Aliasing -> InferredType
Ascribed TypeBase (DimDecl VName) Aliasing
t') ((PatBase Info VName
  -> TermTypeM
       (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
        Occurrences))
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$ \PatBase Info VName
xpat' ->
                  TermTypeM
  (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
   Occurrences)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM a
noUnique (TermTypeM
   (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
    Occurrences)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> ((PatBase Info VName
     -> TermTypeM
          (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
           Occurrences))
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a.
[SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName -> TermTypeM a)
-> TermTypeM a
bindingPat [] UncheckedPat
mergepat (TypeBase (DimDecl VName) Aliasing -> InferredType
Ascribed TypeBase (DimDecl VName) Aliasing
merge_t) ((PatBase Info VName
  -> TermTypeM
       (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
        Occurrences))
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$
                    \PatBase Info VName
mergepat' -> TermTypeM
  (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
   Occurrences)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM a
onlySelfAliasing (TermTypeM
   (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
    Occurrences)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (TermTypeM
      ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeM
  ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM (a, Occurrences)
tapOccurrences (TermTypeM
   ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$ do
                      Exp
loopbody' <- TermTypeM Exp -> TermTypeM Exp
forall a. TermTypeM a -> TermTypeM a
noSizeEscape (TermTypeM Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
loopbody
                      ([VName]
sparams, PatBase Info VName
mergepat'') <- PatBase Info VName
-> Exp -> TermTypeM ([VName], PatBase Info VName)
checkLoopReturnSize PatBase Info VName
mergepat' Exp
loopbody'
                      ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
forall (m :: * -> *) a. Monad m => a -> m a
return
                        ( [VName]
sparams,
                          PatBase Info VName
mergepat'',
                          PatBase Info VName -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn PatBase Info VName
xpat' Exp
e',
                          Exp
loopbody'
                        )
              | Bool
otherwise ->
                SrcLoc
-> Notes
-> Doc
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError (UncheckedExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf UncheckedExp
e) Notes
forall a. Monoid a => a
mempty (Doc
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> Doc
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
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
<+> TypeBase (DimDecl VName) Aliasing -> Doc
forall a. Pretty a => a -> Doc
ppr TypeBase (DimDecl VName) Aliasing
t
        While UncheckedExp
cond ->
          TermTypeM
  (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
   Occurrences)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM a
noUnique (TermTypeM
   (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
    Occurrences)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> ((PatBase Info VName
     -> TermTypeM
          (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
           Occurrences))
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a.
[SizeBinder VName]
-> UncheckedPat
-> InferredType
-> (PatBase Info VName -> TermTypeM a)
-> TermTypeM a
bindingPat [] UncheckedPat
mergepat (TypeBase (DimDecl VName) Aliasing -> InferredType
Ascribed TypeBase (DimDecl VName) Aliasing
merge_t) ((PatBase Info VName
  -> TermTypeM
       (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
        Occurrences))
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (PatBase Info VName
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$ \PatBase Info VName
mergepat' ->
            TermTypeM
  (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
   Occurrences)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM a
onlySelfAliasing (TermTypeM
   (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
    Occurrences)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> ((Exp
     -> Occurrences
     -> TermTypeM
          ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
    -> TermTypeM
         (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
          Occurrences))
-> (Exp
    -> Occurrences
    -> TermTypeM
         ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeM
  ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a. TermTypeM a -> TermTypeM (a, Occurrences)
tapOccurrences
              (TermTypeM
   ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> ((Exp
     -> Occurrences
     -> TermTypeM
          ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
    -> TermTypeM
         ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
-> (Exp
    -> Occurrences
    -> TermTypeM
         ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeM Exp
-> (Exp
    -> Occurrences
    -> TermTypeM
         ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
forall a b.
TermTypeM a -> (a -> Occurrences -> TermTypeM b) -> TermTypeM b
sequentially
                ( UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
cond
                    TermTypeM Exp -> (Exp -> TermTypeM Exp) -> TermTypeM Exp
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> StructType -> Exp -> TermTypeM Exp
unifies String
"being the condition of a 'while' loop" (ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> StructType)
-> ScalarTypeBase (DimDecl VName) () -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
Bool)
                )
              ((Exp
  -> Occurrences
  -> TermTypeM
       ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
 -> TermTypeM
      (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
       Occurrences))
-> (Exp
    -> Occurrences
    -> TermTypeM
         ([VName], PatBase Info VName, LoopFormBase Info VName, Exp))
-> TermTypeM
     (([VName], PatBase Info VName, LoopFormBase Info VName, Exp),
      Occurrences)
forall a b. (a -> b) -> a -> b
$ \Exp
cond' Occurrences
_ -> do
                Exp
loopbody' <- TermTypeM Exp -> TermTypeM Exp
forall a. TermTypeM a -> TermTypeM a
noSizeEscape (TermTypeM Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ UncheckedExp -> TermTypeM Exp
checkExp UncheckedExp
loopbody
                ([VName]
sparams, PatBase Info VName
mergepat'') <- PatBase Info VName
-> Exp -> TermTypeM ([VName], PatBase Info VName)
checkLoopReturnSize PatBase Info VName
mergepat' Exp
loopbody'
                ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
-> TermTypeM
     ([VName], PatBase Info VName, LoopFormBase Info VName, Exp)
forall (m :: * -> *) a. Monad m => a -> m a
return
                  ( [VName]
sparams,
                    PatBase Info VName
mergepat'',
                    Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While Exp
cond',
                    Exp
loopbody'
                  )

    PatBase Info VName
mergepat'' <- do
      TypeBase (DimDecl VName) Aliasing
loopbody_t <- Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
loopbody'
      PatBase Info VName
-> Set VName
-> TypeBase (DimDecl VName) Aliasing
-> Usage
-> TermTypeM (PatBase Info VName)
forall (m :: * -> *) t.
(MonadUnify m, MonadTypeChecker m, Located t,
 MonadReader TermEnv m) =>
PatBase Info VName
-> Set VName
-> TypeBase (DimDecl VName) Aliasing
-> t
-> m (PatBase Info VName)
convergePat PatBase Info VName
mergepat' (Occurrences -> Set VName
allConsumed Occurrences
bodyflow) TypeBase (DimDecl VName) Aliasing
loopbody_t (Usage -> TermTypeM (PatBase Info VName))
-> Usage -> TermTypeM (PatBase Info VName)
forall a b. (a -> b) -> a -> b
$
        SrcLoc -> String -> Usage
mkUsage (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
loopbody') String
"being (part of) the result of the loop body"

    let consumeMerge :: PatBase Info vn -> TypeBase dim Aliasing -> TermTypeM ()
consumeMerge (Id vn
_ (Info TypeBase (DimDecl VName) Aliasing
pt) SrcLoc
ploc) TypeBase dim Aliasing
mt
          | TypeBase (DimDecl VName) Aliasing -> Bool
forall shape as. TypeBase shape as -> Bool
unique TypeBase (DimDecl VName) Aliasing
pt = SrcLoc -> Aliasing -> TermTypeM ()
consume SrcLoc
ploc (Aliasing -> TermTypeM ()) -> Aliasing -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
mt
        consumeMerge (TuplePat [PatBase Info vn]
pats SrcLoc
_) TypeBase dim Aliasing
t
          | Just [TypeBase dim Aliasing]
ts <- TypeBase dim Aliasing -> Maybe [TypeBase dim Aliasing]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord TypeBase dim Aliasing
t =
            (PatBase Info vn -> TypeBase dim Aliasing -> TermTypeM ())
-> [PatBase Info vn] -> [TypeBase dim Aliasing] -> TermTypeM ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ PatBase Info vn -> TypeBase dim Aliasing -> TermTypeM ()
consumeMerge [PatBase Info vn]
pats [TypeBase dim Aliasing]
ts
        consumeMerge (PatParens PatBase Info vn
pat SrcLoc
_) TypeBase dim Aliasing
t =
          PatBase Info vn -> TypeBase dim Aliasing -> TermTypeM ()
consumeMerge PatBase Info vn
pat TypeBase dim Aliasing
t
        consumeMerge (PatAscription PatBase Info vn
pat TypeDeclBase Info vn
_ SrcLoc
_) TypeBase dim Aliasing
t =
          PatBase Info vn -> TypeBase dim Aliasing -> TermTypeM ()
consumeMerge PatBase Info vn
pat TypeBase dim Aliasing
t
        consumeMerge PatBase Info vn
_ TypeBase dim Aliasing
_ =
          () -> TermTypeM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    PatBase Info VName
-> TypeBase (DimDecl VName) Aliasing -> TermTypeM ()
forall vn dim.
PatBase Info vn -> TypeBase dim Aliasing -> TermTypeM ()
consumeMerge PatBase Info VName
mergepat'' (TypeBase (DimDecl VName) Aliasing -> TermTypeM ())
-> TermTypeM (TypeBase (DimDecl VName) Aliasing) -> TermTypeM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
mergeexp'

    -- dim handling (3)
    StructType
merge_t' <-
      SrcLoc
-> Rigidity
-> Name
-> Set VName
-> StructType
-> TermTypeM StructType
forall als.
SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) als
-> TermTypeM (TypeBase (DimDecl VName) als)
someDimsFreshInType SrcLoc
loc Rigidity
Nonrigid Name
"loop" ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams) (StructType -> TermTypeM StructType)
-> StructType -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$
        TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (TypeBase (DimDecl VName) Aliasing -> StructType)
-> TypeBase (DimDecl VName) Aliasing -> StructType
forall a b. (a -> b) -> a -> b
$ PatBase Info VName -> TypeBase (DimDecl VName) Aliasing
patternType PatBase Info VName
mergepat''
    StructType
mergeexp_t <- TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (TypeBase (DimDecl VName) Aliasing -> StructType)
-> TermTypeM (TypeBase (DimDecl VName) Aliasing)
-> TermTypeM StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> TermTypeM (TypeBase (DimDecl VName) Aliasing)
expTypeFully Exp
mergeexp'
    Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (StructType -> StructType -> Checking
CheckingLoopInitial StructType
merge_t' StructType
mergeexp_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
        (SrcLoc -> String -> Usage
mkUsage (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
mergeexp') String
"matching initial loop values to pattern")
        StructType
merge_t'
        StructType
mergeexp_t

    (TypeBase (DimDecl VName) Aliasing
loopt, [VName]
retext) <-
      SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) Aliasing
-> TermTypeM (TypeBase (DimDecl VName) Aliasing, [VName])
forall als.
SrcLoc
-> Rigidity
-> Name
-> Set VName
-> TypeBase (DimDecl VName) als
-> TermTypeM (TypeBase (DimDecl VName) als, [VName])
freshDimsInType SrcLoc
loc (RigidSource -> Rigidity
Rigid RigidSource
RigidLoop) Name
"loop" ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams) (TypeBase (DimDecl VName) Aliasing
 -> TermTypeM (TypeBase (DimDecl VName) Aliasing, [VName]))
-> TypeBase (DimDecl VName) Aliasing
-> TermTypeM (TypeBase (DimDecl VName) Aliasing, [VName])
forall a b. (a -> b) -> a -> b
$
        PatBase Info VName -> TypeBase (DimDecl VName) Aliasing
patternType PatBase Info VName
mergepat''
    -- We set all of the uniqueness to be unique.  This is intentional,
    -- and matches what happens for function calls.  Those arrays that
    -- really *cannot* be consumed will alias something unconsumable,
    -- and will be caught that way.
    let bound_here :: Set VName
bound_here = PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
mergepat'' 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]
sparams Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
form_bound
        form_bound :: Set VName
form_bound =
          case LoopFormBase Info VName
form' of
            For Ident
v Exp
_ -> VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ Ident -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName Ident
v
            ForIn PatBase Info VName
forpat Exp
_ -> PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
forpat
            While {} -> Set VName
forall a. Monoid a => a
mempty
        loopt' :: TypeBase (DimDecl VName) Aliasing
loopt' =
          (Aliasing -> Aliasing)
-> TypeBase (DimDecl VName) Aliasing
-> TypeBase (DimDecl VName) Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliasing -> Aliasing -> Aliasing
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` (VName -> Alias) -> Set VName -> Aliasing
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map VName -> Alias
AliasBound Set VName
bound_here) (TypeBase (DimDecl VName) Aliasing
 -> TypeBase (DimDecl VName) Aliasing)
-> TypeBase (DimDecl VName) Aliasing
-> TypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$
            TypeBase (DimDecl VName) Aliasing
loopt TypeBase (DimDecl VName) Aliasing
-> Uniqueness -> TypeBase (DimDecl VName) Aliasing
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Unique

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

    (CheckedLoop, AppRes) -> TermTypeM (CheckedLoop, AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([VName]
sparams, PatBase Info VName
mergepat'', Exp
mergeexp', LoopFormBase Info VName
form', Exp
loopbody'), TypeBase (DimDecl VName) Aliasing -> [VName] -> AppRes
AppRes TypeBase (DimDecl VName) Aliasing
loopt' [VName]
retext)
  where
    convergePat :: PatBase Info VName
-> Set VName
-> TypeBase (DimDecl VName) Aliasing
-> t
-> m (PatBase Info VName)
convergePat PatBase Info VName
pat Set VName
body_cons TypeBase (DimDecl VName) Aliasing
body_t t
body_loc = do
      let consumed_merge :: Set VName
consumed_merge = PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
pat Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set VName
body_cons

          uniquePat :: PatBase Info VName -> PatBase Info VName
uniquePat (Wildcard (Info TypeBase (DimDecl VName) Aliasing
t) SrcLoc
wloc) =
            Info (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn.
f (TypeBase (DimDecl VName) Aliasing) -> SrcLoc -> PatBase f vn
Wildcard (TypeBase (DimDecl VName) Aliasing
-> Info (TypeBase (DimDecl VName) Aliasing)
forall a. a -> Info a
Info (TypeBase (DimDecl VName) Aliasing
 -> Info (TypeBase (DimDecl VName) Aliasing))
-> TypeBase (DimDecl VName) Aliasing
-> Info (TypeBase (DimDecl VName) Aliasing)
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) Aliasing
t TypeBase (DimDecl VName) Aliasing
-> Uniqueness -> TypeBase (DimDecl VName) Aliasing
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique) SrcLoc
wloc
          uniquePat (PatParens PatBase Info VName
p SrcLoc
ploc) =
            PatBase Info VName -> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens (PatBase Info VName -> PatBase Info VName
uniquePat PatBase Info VName
p) SrcLoc
ploc
          uniquePat (PatAttr AttrInfo VName
attr PatBase Info VName
p SrcLoc
ploc) =
            AttrInfo VName
-> PatBase Info VName -> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr (PatBase Info VName -> PatBase Info VName
uniquePat PatBase Info VName
p) SrcLoc
ploc
          uniquePat (Id VName
name (Info TypeBase (DimDecl VName) Aliasing
t) SrcLoc
iloc)
            | VName
name VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
consumed_merge =
              let t' :: TypeBase (DimDecl VName) Aliasing
t' = TypeBase (DimDecl VName) Aliasing
t TypeBase (DimDecl VName) Aliasing
-> Uniqueness -> TypeBase (DimDecl VName) Aliasing
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Unique TypeBase (DimDecl VName) Aliasing
-> Aliasing -> TypeBase (DimDecl VName) Aliasing
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliasing
forall a. Monoid a => a
mempty
               in VName
-> Info (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc
-> PatBase Info VName
forall (f :: * -> *) vn.
vn
-> f (TypeBase (DimDecl VName) Aliasing) -> SrcLoc -> PatBase f vn
Id VName
name (TypeBase (DimDecl VName) Aliasing
-> Info (TypeBase (DimDecl VName) Aliasing)
forall a. a -> Info a
Info TypeBase (DimDecl VName) Aliasing
t') SrcLoc
iloc
            | Bool
otherwise =
              let t' :: TypeBase (DimDecl VName) Aliasing
t' = TypeBase (DimDecl VName) Aliasing
t TypeBase (DimDecl VName) Aliasing
-> Uniqueness -> TypeBase (DimDecl VName) Aliasing
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
               in VName
-> Info (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc
-> PatBase Info VName
forall (f :: * -> *) vn.
vn
-> f (TypeBase (DimDecl VName) Aliasing) -> SrcLoc -> PatBase f vn
Id VName
name (TypeBase (DimDecl VName) Aliasing
-> Info (TypeBase (DimDecl VName) Aliasing)
forall a. a -> Info a
Info TypeBase (DimDecl VName) Aliasing
t') SrcLoc
iloc
          uniquePat (TuplePat [PatBase Info VName]
pats SrcLoc
ploc) =
            [PatBase Info VName] -> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat ((PatBase Info VName -> PatBase Info VName)
-> [PatBase Info VName] -> [PatBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map PatBase Info VName -> PatBase Info VName
uniquePat [PatBase Info VName]
pats) SrcLoc
ploc
          uniquePat (RecordPat [(Name, PatBase Info VName)]
fs SrcLoc
ploc) =
            [(Name, PatBase Info VName)] -> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat (((Name, PatBase Info VName) -> (Name, PatBase Info VName))
-> [(Name, PatBase Info VName)] -> [(Name, PatBase Info VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((PatBase Info VName -> PatBase Info VName)
-> (Name, PatBase Info VName) -> (Name, PatBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PatBase Info VName -> PatBase Info VName
uniquePat) [(Name, PatBase Info VName)]
fs) SrcLoc
ploc
          uniquePat (PatAscription PatBase Info VName
p TypeDeclBase Info VName
t SrcLoc
ploc) =
            PatBase Info VName
-> TypeDeclBase Info VName -> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn.
PatBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatBase f vn
PatAscription PatBase Info VName
p TypeDeclBase Info VName
t SrcLoc
ploc
          uniquePat p :: PatBase Info VName
p@PatLit {} = PatBase Info VName
p
          uniquePat (PatConstr Name
n Info (TypeBase (DimDecl VName) Aliasing)
t [PatBase Info VName]
ps SrcLoc
ploc) =
            Name
-> Info (TypeBase (DimDecl VName) Aliasing)
-> [PatBase Info VName]
-> SrcLoc
-> PatBase Info VName
forall (f :: * -> *) vn.
Name
-> f (TypeBase (DimDecl VName) Aliasing)
-> [PatBase f vn]
-> SrcLoc
-> PatBase f vn
PatConstr Name
n Info (TypeBase (DimDecl VName) Aliasing)
t ((PatBase Info VName -> PatBase Info VName)
-> [PatBase Info VName] -> [PatBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map PatBase Info VName -> PatBase Info VName
uniquePat [PatBase Info VName]
ps) SrcLoc
ploc

          -- Make the pattern unique where needed.
          pat' :: PatBase Info VName
pat' = PatBase Info VName -> PatBase Info VName
uniquePat PatBase Info VName
pat

      TypeBase (DimDecl VName) Aliasing
pat_t <- TypeBase (DimDecl VName) Aliasing
-> m (TypeBase (DimDecl VName) Aliasing)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully (TypeBase (DimDecl VName) Aliasing
 -> m (TypeBase (DimDecl VName) Aliasing))
-> TypeBase (DimDecl VName) Aliasing
-> m (TypeBase (DimDecl VName) Aliasing)
forall a b. (a -> b) -> a -> b
$ PatBase Info VName -> TypeBase (DimDecl VName) Aliasing
patternType PatBase Info VName
pat'
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (TypeBase (DimDecl VName) Aliasing -> TypeBase () ()
forall dim as. TypeBase dim as -> TypeBase () ()
toStructural TypeBase (DimDecl VName) Aliasing
body_t TypeBase () () -> TypeBase () () -> Bool
`subtypeOf` TypeBase (DimDecl VName) Aliasing -> TypeBase () ()
forall dim as. TypeBase dim as -> TypeBase () ()
toStructural TypeBase (DimDecl VName) Aliasing
pat_t) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
        SrcLoc -> StructType -> [StructType] -> m ()
forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> StructType -> [StructType] -> m a
unexpectedType (t -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf t
body_loc) (TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
body_t) [TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
pat_t]

      -- Check that the new values of consumed merge parameters do not
      -- alias something bound outside the loop, AND that anything
      -- returned for a unique merge parameter does not alias anything
      -- else returned.  We also update the aliases for the pattern.
      Set VName
bound_outside <- (TermEnv -> Set VName) -> m (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((TermEnv -> Set VName) -> m (Set VName))
-> (TermEnv -> Set VName) -> m (Set VName)
forall a b. (a -> b) -> a -> b
$ [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName)
-> (TermEnv -> [VName]) -> TermEnv -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName ValBinding -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName ValBinding -> [VName])
-> (TermEnv -> Map VName ValBinding) -> TermEnv -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermScope -> Map VName ValBinding
scopeVtable (TermScope -> Map VName ValBinding)
-> (TermEnv -> TermScope) -> TermEnv -> Map VName ValBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermEnv -> TermScope
termScope
      let combAliases :: TypeBase dim ast -> TypeBase shape ast -> TypeBase dim ast
combAliases TypeBase dim ast
t1 TypeBase shape ast
t2 =
            case TypeBase dim ast
t1 of
              Scalar Record {} -> TypeBase dim ast
t1
              TypeBase dim ast
_ -> TypeBase dim ast
t1 TypeBase dim ast -> (ast -> ast) -> TypeBase dim ast
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` (ast -> ast -> ast
forall a. Semigroup a => a -> a -> a
<> TypeBase shape ast -> ast
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase shape ast
t2)

          checkMergeReturn :: PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn)
checkMergeReturn (Id vn
pat_v (Info TypeBase (DimDecl VName) Aliasing
pat_v_t) SrcLoc
patloc) TypeBase dim Aliasing
t
            | TypeBase (DimDecl VName) Aliasing -> Bool
forall shape as. TypeBase shape as -> Bool
unique TypeBase (DimDecl VName) Aliasing
pat_v_t,
              VName
v : [VName]
_ <-
                Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [VName]
forall a b. (a -> b) -> a -> b
$
                  (Alias -> VName) -> Aliasing -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar (TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t) Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set VName
bound_outside =
              m (PatBase Info vn) -> t m (PatBase Info vn)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (PatBase Info vn) -> t m (PatBase Info vn))
-> (Doc -> m (PatBase Info vn)) -> Doc -> t m (PatBase Info vn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Notes -> Doc -> m (PatBase Info vn)
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> t m (PatBase Info vn)) -> Doc -> t m (PatBase Info vn)
forall a b. (a -> b) -> a -> b
$
                Doc
"Return value for loop parameter"
                  Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (vn -> Doc
forall v. IsName v => v -> Doc
pprName vn
pat_v)
                  Doc -> Doc -> Doc
<+> Doc
"aliases"
                  Doc -> Doc -> Doc
<+> VName -> Doc
forall v. IsName v => v -> Doc
pprName VName
v Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."
            | Bool
otherwise = do
              (Aliasing
cons, Aliasing
obs) <- t m (Aliasing, Aliasing)
forall s (m :: * -> *). MonadState s m => m s
get
              Bool -> t m () -> t m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Aliasing -> Bool
forall a. Set a -> Bool
S.null (Aliasing -> Bool) -> Aliasing -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t Aliasing -> Aliasing -> Aliasing
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Aliasing
cons) (t m () -> t m ()) -> t m () -> t m ()
forall a b. (a -> b) -> a -> b
$
                m () -> t m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> t m ()) -> (Doc -> m ()) -> Doc -> t m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Notes -> Doc -> m ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> t m ()) -> Doc -> t m ()
forall a b. (a -> b) -> a -> b
$
                  Doc
"Return value for loop parameter"
                    Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (vn -> Doc
forall v. IsName v => v -> Doc
pprName vn
pat_v)
                    Doc -> Doc -> Doc
<+> Doc
"aliases other consumed loop parameter."
              Bool -> t m () -> t m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
                ( TypeBase (DimDecl VName) Aliasing -> Bool
forall shape as. TypeBase shape as -> Bool
unique TypeBase (DimDecl VName) Aliasing
pat_v_t
                    Bool -> Bool -> Bool
&& Bool -> Bool
not (Aliasing -> Bool
forall a. Set a -> Bool
S.null (TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t Aliasing -> Aliasing -> Aliasing
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (Aliasing
cons Aliasing -> Aliasing -> Aliasing
forall a. Semigroup a => a -> a -> a
<> Aliasing
obs)))
                )
                (t m () -> t m ()) -> t m () -> t m ()
forall a b. (a -> b) -> a -> b
$ m () -> t m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> t m ()) -> (Doc -> m ()) -> Doc -> t m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Notes -> Doc -> m ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> t m ()) -> Doc -> t m ()
forall a b. (a -> b) -> a -> b
$
                  Doc
"Return value for consuming loop parameter"
                    Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (vn -> Doc
forall v. IsName v => v -> Doc
pprName vn
pat_v)
                    Doc -> Doc -> Doc
<+> Doc
"aliases previously returned value."
              if TypeBase (DimDecl VName) Aliasing -> Bool
forall shape as. TypeBase shape as -> Bool
unique TypeBase (DimDecl VName) Aliasing
pat_v_t
                then (Aliasing, Aliasing) -> t m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Aliasing
cons Aliasing -> Aliasing -> Aliasing
forall a. Semigroup a => a -> a -> a
<> TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t, Aliasing
obs)
                else (Aliasing, Aliasing) -> t m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Aliasing
cons, Aliasing
obs Aliasing -> Aliasing -> Aliasing
forall a. Semigroup a => a -> a -> a
<> TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t)

              PatBase Info vn -> t m (PatBase Info vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatBase Info vn -> t m (PatBase Info vn))
-> PatBase Info vn -> t m (PatBase Info vn)
forall a b. (a -> b) -> a -> b
$ vn
-> Info (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc
-> PatBase Info vn
forall (f :: * -> *) vn.
vn
-> f (TypeBase (DimDecl VName) Aliasing) -> SrcLoc -> PatBase f vn
Id vn
pat_v (TypeBase (DimDecl VName) Aliasing
-> Info (TypeBase (DimDecl VName) Aliasing)
forall a. a -> Info a
Info (TypeBase (DimDecl VName) Aliasing
-> TypeBase dim Aliasing -> TypeBase (DimDecl VName) Aliasing
forall ast dim shape.
Monoid ast =>
TypeBase dim ast -> TypeBase shape ast -> TypeBase dim ast
combAliases TypeBase (DimDecl VName) Aliasing
pat_v_t TypeBase dim Aliasing
t)) SrcLoc
patloc
          checkMergeReturn (Wildcard (Info TypeBase (DimDecl VName) Aliasing
pat_v_t) SrcLoc
patloc) TypeBase dim Aliasing
t =
            PatBase Info vn -> t m (PatBase Info vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatBase Info vn -> t m (PatBase Info vn))
-> PatBase Info vn -> t m (PatBase Info vn)
forall a b. (a -> b) -> a -> b
$ Info (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn.
f (TypeBase (DimDecl VName) Aliasing) -> SrcLoc -> PatBase f vn
Wildcard (TypeBase (DimDecl VName) Aliasing
-> Info (TypeBase (DimDecl VName) Aliasing)
forall a. a -> Info a
Info (TypeBase (DimDecl VName) Aliasing
-> TypeBase dim Aliasing -> TypeBase (DimDecl VName) Aliasing
forall ast dim shape.
Monoid ast =>
TypeBase dim ast -> TypeBase shape ast -> TypeBase dim ast
combAliases TypeBase (DimDecl VName) Aliasing
pat_v_t TypeBase dim Aliasing
t)) SrcLoc
patloc
          checkMergeReturn (PatParens PatBase Info vn
p SrcLoc
_) TypeBase dim Aliasing
t =
            PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn)
checkMergeReturn PatBase Info vn
p TypeBase dim Aliasing
t
          checkMergeReturn (PatAscription PatBase Info vn
p TypeDeclBase Info vn
_ SrcLoc
_) TypeBase dim Aliasing
t =
            PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn)
checkMergeReturn PatBase Info vn
p TypeBase dim Aliasing
t
          checkMergeReturn (RecordPat [(Name, PatBase Info vn)]
pfs SrcLoc
patloc) (Scalar (Record Map Name (TypeBase dim Aliasing)
tfs)) =
            [(Name, PatBase Info vn)] -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat ([(Name, PatBase Info vn)] -> SrcLoc -> PatBase Info vn)
-> (Map Name (PatBase Info vn) -> [(Name, PatBase Info vn)])
-> Map Name (PatBase Info vn)
-> SrcLoc
-> PatBase Info vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (PatBase Info vn) -> [(Name, PatBase Info vn)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name (PatBase Info vn) -> SrcLoc -> PatBase Info vn)
-> t m (Map Name (PatBase Info vn))
-> t m (SrcLoc -> PatBase Info vn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (t m (PatBase Info vn))
-> t m (Map Name (PatBase Info vn))
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence Map Name (t m (PatBase Info vn))
pfs' t m (SrcLoc -> PatBase Info vn)
-> t m SrcLoc -> t m (PatBase Info vn)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
patloc
            where
              pfs' :: Map Name (t m (PatBase Info vn))
pfs' = (PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn))
-> Map Name (PatBase Info vn)
-> Map Name (TypeBase dim Aliasing)
-> Map Name (t m (PatBase Info vn))
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn)
checkMergeReturn ([(Name, PatBase Info vn)] -> Map Name (PatBase Info vn)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatBase Info vn)]
pfs) Map Name (TypeBase dim Aliasing)
tfs
          checkMergeReturn (TuplePat [PatBase Info vn]
pats SrcLoc
patloc) TypeBase dim Aliasing
t
            | Just [TypeBase dim Aliasing]
ts <- TypeBase dim Aliasing -> Maybe [TypeBase dim Aliasing]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord TypeBase dim Aliasing
t =
              [PatBase Info vn] -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat ([PatBase Info vn] -> SrcLoc -> PatBase Info vn)
-> t m [PatBase Info vn] -> t m (SrcLoc -> PatBase Info vn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn))
-> [PatBase Info vn]
-> [TypeBase dim Aliasing]
-> t m [PatBase Info vn]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn)
checkMergeReturn [PatBase Info vn]
pats [TypeBase dim Aliasing]
ts t m (SrcLoc -> PatBase Info vn)
-> t m SrcLoc -> t m (PatBase Info vn)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
patloc
          checkMergeReturn PatBase Info vn
p TypeBase dim Aliasing
_ =
            PatBase Info vn -> t m (PatBase Info vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure PatBase Info vn
p

      (PatBase Info VName
pat'', (Aliasing
pat_cons, Aliasing
_)) <-
        StateT (Aliasing, Aliasing) m (PatBase Info VName)
-> (Aliasing, Aliasing)
-> m (PatBase Info VName, (Aliasing, Aliasing))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (PatBase Info VName
-> TypeBase (DimDecl VName) Aliasing
-> StateT (Aliasing, Aliasing) m (PatBase Info VName)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) vn dim.
(MonadTrans t, MonadTypeChecker m, IsName vn,
 MonadState (Aliasing, Aliasing) (t m)) =>
PatBase Info vn -> TypeBase dim Aliasing -> t m (PatBase Info vn)
checkMergeReturn PatBase Info VName
pat' TypeBase (DimDecl VName) Aliasing
body_t) (Aliasing
forall a. Monoid a => a
mempty, Aliasing
forall a. Monoid a => a
mempty)

      let body_cons' :: Set VName
body_cons' = Set VName
body_cons Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (Alias -> VName) -> Aliasing -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar Aliasing
pat_cons
      if Set VName
body_cons' Set VName -> Set VName -> Bool
forall a. Eq a => a -> a -> Bool
== Set VName
body_cons Bool -> Bool -> Bool
&& PatBase Info VName -> TypeBase (DimDecl VName) Aliasing
patternType PatBase Info VName
pat'' TypeBase (DimDecl VName) Aliasing
-> TypeBase (DimDecl VName) Aliasing -> Bool
forall a. Eq a => a -> a -> Bool
== PatBase Info VName -> TypeBase (DimDecl VName) Aliasing
patternType PatBase Info VName
pat
        then PatBase Info VName -> m (PatBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure PatBase Info VName
pat'
        else PatBase Info VName
-> Set VName
-> TypeBase (DimDecl VName) Aliasing
-> t
-> m (PatBase Info VName)
convergePat PatBase Info VName
pat'' Set VName
body_cons' TypeBase (DimDecl VName) Aliasing
body_t t
body_loc