-- | Size inference of @loop@. This is complicated, 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 o ->
  TermTypeM (TypeBase Size o)
someDimsFreshInType :: forall o.
SrcLoc
-> Name -> [VName] -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
someDimsFreshInType SrcLoc
loc Name
desc [VName]
dims TypeBase Exp o
t = do
  areSameSize <- TermTypeM (VName -> VName -> Bool)
forall (m :: * -> *). MonadUnify m => m (VName -> VName -> Bool)
getAreSame
  let 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]
dims
  bitraverse (onDim freshen) pure 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
          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
          pure $ sizeFromName (qualName v) 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 o ->
  TermTypeM (TypeBase Size o, [VName])
freshDimsInType :: forall o.
Usage
-> Rigidity
-> Name
-> [VName]
-> TypeBase Exp o
-> TermTypeM (TypeBase Exp o, [VName])
freshDimsInType Usage
usage Rigidity
r Name
desc [VName]
dims TypeBase Exp o
t = do
  areSameSize <- TermTypeM (VName -> VName -> Bool)
forall (m :: * -> *). MonadUnify m => m (VName -> VName -> Bool)
getAreSame
  second (map snd) <$> runStateT (bitraverse (onDim areSameSize) pure t) 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]
dims = do
          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 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
              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
              modify ((d, v) :)
              pure $ sizeFromName (qualName v) $ srclocOf 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

-- | How the loop body produces the size at a given loop parameter position,
-- relative to the fresh variable standing for that position's size. See Note
-- [Loop size inference].
data DimClass
  = -- | A genuinely new size: this position is variant.
    New
  | -- | The body produces this position's initial size (the size the fresh
    -- variable replaced). Either it produces that size directly (@Nothing@), or
    -- it copies the /current/ size of another parameter @u@ that started at the
    -- same initial size (@Just u@). In the latter case the two sizes coincide
    -- only while @u@ is unchanged, so this position is invariant exactly when
    -- @u@ is.
    Reproduces (Maybe VName)

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
  (loop_t, _) <-
    Usage
-> Rigidity
-> Name
-> [VName]
-> StructType
-> TermTypeM (StructType, [VName])
forall o.
Usage
-> Rigidity
-> Name
-> [VName]
-> TypeBase Exp o
-> TermTypeM (TypeBase Exp o, [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 o. TypeBase dim o -> TypeBase dim NoMode
toStruct (Pat ParamType -> ParamType
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType Pat ParamType
pat)
  arg_t <- toStruct <$> expTypeFully arg
  onFailure (checking loop_t arg_t) $
    unify (mkUsage arg desc) loop_t 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 =
  (Pat ParamType, LoopInitBase Info VName, LoopFormBase Info VName, Exp)

-- | 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
  cs <- TermTypeM (Map VName (Level, Constraint))
forall (m :: * -> *).
MonadUnify m =>
m (Map VName (Level, Constraint))
getConstraints
  let 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
        (_, UnknownSize v_loc _) <- 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
        Just . typeError (srclocOf loc) mempty $
          "Inferred type for loop parameter is"
            </> indent 2 (pretty pat_t)
            </> "but"
            <+> dquotes (prettyName v)
            <+> "is an existential size created inside the loop body at"
            <+> pretty (locStrRel loc v_loc)
            <> "."
  case mapMaybe bad $ S.toList $ fvVars $ freeInType pat_t of
    TermTypeM ()
problem : [TermTypeM ()]
_ -> TermTypeM ()
problem
    [] -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- Check if any sparams have reduced levels, which indicates they escaped the
-- loop scope and got constrained by something at a lower (outer) level. This
-- means they cannot actually be variant across the loop.
checkForEscaped :: SrcLoc -> M.Map VName Level -> [VName] -> TermTypeM ()
checkForEscaped :: SrcLoc -> Map VName Level -> [VName] -> TermTypeM ()
checkForEscaped SrcLoc
loc Map VName Level
initial_levels [VName]
sparams = do
  cs <- TermTypeM (Map VName (Level, Constraint))
forall (m :: * -> *).
MonadUnify m =>
m (Map VName (Level, Constraint))
getConstraints
  let checkSparamLevel VName
sparam = do
        let current_level :: Maybe Level
current_level = (Level, Constraint) -> Level
forall a b. (a, b) -> a
fst ((Level, Constraint) -> Level)
-> Maybe (Level, Constraint) -> Maybe Level
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (Level, Constraint) -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
sparam Map VName (Level, Constraint)
cs
            original_level :: Maybe Level
original_level = VName -> Map VName Level -> Maybe Level
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
sparam Map VName Level
initial_levels
        case (Maybe Level
original_level, Maybe Level
current_level) of
          (Just Level
orig, Just Level
curr)
            | Level
curr Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
< Level
orig ->
                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 SrcLoc
loc Notes
forall a. Monoid a => a
mempty
                  (Doc () -> m a) -> (Doc () -> Doc ()) -> Doc () -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"loop-variant-escape"
                  (Doc () -> Maybe (m a)) -> Doc () -> Maybe (m a)
forall a b. (a -> b) -> a -> b
$ Doc ()
"Loop-variant size"
                    Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"has escaped into type that occurs outside of loop."
                    Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"This is likely because the loop body uses a variable whose type"
                    Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"depends on the loop parameter."
          (Maybe Level, Maybe Level)
_ -> Maybe (m a)
forall a. Maybe a
Nothing
  case mapMaybe checkSparamLevel sparams 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.  See Note [Loop size inference].
checkLoop ::
  (Exp -> TermTypeM Exp) ->
  UncheckedLoop ->
  SrcLoc ->
  TermTypeM (CheckedLoop, AppRes)
checkLoop :: (Exp -> TermTypeM Exp)
-> UncheckedLoop -> SrcLoc -> TermTypeM (CheckedLoop, AppRes)
checkLoop Exp -> TermTypeM Exp
checkExp (Pat ParamType
looppat, LoopInitBase Info VName
loopinit, LoopFormBase Info VName
form, Exp
loopbody) SrcLoc
loc = do
  loopinit' <- Exp -> TermTypeM Exp
checkExp (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
loopinit
  known_before <- M.keysSet <$> getConstraints

  (loop_t, new_dims_map) <-
    -- dim handling (1)
    allDimsFreshInType
      (mkUsage loc "loop parameter type inference")
      Nonrigid
      "loop_d"
      =<< expTypeFully loopinit'
  let 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, 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

  -- Save the initial levels of new_dims for later checking
  initial_levels <- do
    cs <- getConstraints
    pure $ M.fromList [(v, lvl) | v <- new_dims, Just (lvl, _) <- [M.lookup v cs]]

  -- dim handling (2)
  let checkLoopReturnSize Pat ParamType
looppat' Exp
loopbody' = do
        loopbody_t <- Exp -> TermTypeM StructType
expTypeFully Exp
loopbody'
        looppat_t <- normTypeFully (patternType looppat')

        let 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
        checkForImpossible (locOf looppat) ok_names looppat_t

        pat_t <- someDimsFreshInType loc "loop" new_dims looppat_t

        -- We are ignoring the dimensions here, because any mismatches
        -- should be turned into fresh size variables.
        onFailure (CheckingLoopBody (toStruct pat_t) (toStruct loopbody_t)) $
          unify
            (mkUsage loopbody "matching loop body to loop pattern")
            (toStruct pat_t)
            (toStruct 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.
        areSameSize <- getAreSame
        let initialOf VName
v = (VName, Exp) -> Exp
forall a b. (a, b) -> b
snd ((VName, Exp) -> Exp) -> Maybe (VName, Exp) -> Maybe Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((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
            sameSize (Var QualName VName
x f StructType
_ SrcLoc
_) (Var QualName VName
y f StructType
_ SrcLoc
_) = VName -> VName -> Bool
areSameSize (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
x) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
y)
            sameSize ExpBase f VName
x ExpBase f VName
y = ExpBase f VName
x ExpBase f VName -> ExpBase f VName -> Bool
forall a. Eq a => a -> a -> Bool
== ExpBase f VName
y
            -- If the body-produced size 'd' is another loop parameter that
            -- shares the initial size 'e'', return that parameter.
            sharesInitial Exp
e' (Var QualName VName
d f StructType
_ SrcLoc
_)
              | Just Exp
d_init <- VName -> Maybe Exp
initialOf (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d),
                Exp -> Exp -> Bool
forall {f :: * -> *}.
Eq (ExpBase f VName) =>
ExpBase f VName -> ExpBase f VName -> Bool
sameSize Exp
d_init Exp
e' =
                  VName -> Maybe VName
forall a. a -> Maybe a
Just (VName -> Maybe VName) -> VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
            sharesInitial Exp
_ ExpBase f VName
_ = Maybe VName
forall a. Maybe a
Nothing
            -- Classify each new_dim 'v' (with initial size 'e'') by how the
            -- loop body produces the corresponding size 'd'.
            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 ->
                case VName -> Maybe Exp
initialOf VName
v of
                  Just Exp
e'
                    | Exp -> Exp -> Bool
forall {f :: * -> *}.
Eq (ExpBase f VName) =>
ExpBase f VName -> ExpBase f VName -> Bool
sameSize Exp
e' Exp
d ->
                        (Map VName DimClass -> Map VName DimClass) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map VName DimClass -> Map VName DimClass) -> f ())
-> (Map VName DimClass -> Map VName DimClass) -> f ()
forall a b. (a -> b) -> a -> b
$ VName -> DimClass -> Map VName DimClass -> Map VName DimClass
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (DimClass -> Map VName DimClass -> Map VName DimClass)
-> DimClass -> Map VName DimClass -> Map VName DimClass
forall a b. (a -> b) -> a -> b
$ Maybe VName -> DimClass
Reproduces Maybe VName
forall a. Maybe a
Nothing
                    | VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
new_dims,
                      Just VName
u <- Exp -> Exp -> Maybe VName
forall {f :: * -> *}. Exp -> ExpBase f VName -> Maybe VName
sharesInitial Exp
e' Exp
d ->
                        (Map VName DimClass -> Map VName DimClass) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map VName DimClass -> Map VName DimClass) -> f ())
-> (Map VName DimClass -> Map VName DimClass) -> f ()
forall a b. (a -> b) -> a -> b
$ VName -> DimClass -> Map VName DimClass -> Map VName DimClass
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (DimClass -> Map VName DimClass -> Map VName DimClass)
-> DimClass -> Map VName DimClass -> Map VName DimClass
forall a b. (a -> b) -> a -> b
$ Maybe VName -> DimClass
Reproduces (VName -> Maybe VName
forall a. a -> Maybe a
Just VName
u)
                    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
known_before ->
                        (Map VName DimClass -> Map VName DimClass) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map VName DimClass -> Map VName DimClass) -> f ())
-> (Map VName DimClass -> Map VName DimClass) -> f ()
forall a b. (a -> b) -> a -> b
$ VName -> DimClass -> Map VName DimClass -> Map VName DimClass
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v DimClass
New
                    | Bool
otherwise ->
                        () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                  Maybe Exp
Nothing ->
                    () -> 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
        loopbody_t' <- normTypeFully loopbody_t
        loop_t' <- normTypeFully loop_t

        let classified =
              State (Map VName DimClass) StructType
-> Map VName DimClass -> Map VName DimClass
forall s a. State s a -> s -> s
execState (([VName] -> Exp -> Exp -> StateT (Map VName DimClass) Identity Exp)
-> StructType
-> StructType
-> State (Map VName DimClass) StructType
forall o (m :: * -> *) d1 d2.
(Monoid o, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 o -> TypeBase d2 o -> m (TypeBase d1 o)
matchDims [VName] -> Exp -> Exp -> StateT (Map VName DimClass) Identity Exp
forall {f :: * -> *} {p}.
MonadState (Map VName DimClass) f =>
p -> Exp -> Exp -> f Exp
onDims StructType
loop_t' StructType
loopbody_t') Map VName DimClass
forall a. Monoid a => a
mempty
            -- The variant sizes are the least set containing every 'New'
            -- position and closed under the copies-from dependency.  See Note
            -- [Loop size inference].
            seeds = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName
v | (VName
v, DimClass
New) <- Map VName DimClass -> [(VName, DimClass)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName DimClass
classified]
            grow Set VName
vs =
              Set VName
vs
                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
v
                  | (VName
v, Reproduces (Just VName
u)) <- Map VName DimClass -> [(VName, DimClass)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName DimClass
classified,
                    VName
u VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
vs
                  ]
            fixVariant Set VName
vs =
              let vs' :: Set VName
vs' = Set VName -> Set VName
grow Set VName
vs
               in if Set VName
vs' Set VName -> Set VName -> Bool
forall a. Eq a => a -> a -> Bool
== Set VName
vs then Set VName
vs else Set VName -> Set VName
fixVariant Set VName
vs'
            variant = Set VName -> Set VName
fixVariant Set VName
seeds
            -- The invariant positions are resolved back to the initial size
            -- they stand for.
            init_substs =
              [(VName, Subst t)] -> Map VName (Subst t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Subst t)] -> Map VName (Subst t))
-> [(VName, Subst t)] -> Map VName (Subst t)
forall a b. (a -> b) -> a -> b
$ do
                (v, Reproduces _) <- Map VName DimClass -> [(VName, DimClass)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName DimClass
classified
                guard $ v `S.notMember` variant
                Just e' <- [initialOf v]
                pure (v, ExpSubst e')
            sparams = Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
variant

        checkForEscaped loc initial_levels sparams

        -- 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
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 ()
        mapM_ dimToInit $ M.toList init_substs

        looppat'' <- applySubst (`M.lookup` init_substs) <$> updateTypes looppat'

        -- Eliminate those new_dims that turned into sparams so it won't
        -- look like we have ambiguous sizes lying around.
        modifyConstraints $ M.filterWithKey $ \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.
        wellTypedLoopArg BodyResult sparams looppat'' loopbody'

        pure (nubOrd sparams, looppat'')

  (sparams, looppat', form', loopbody') <-
    case form of
      For IdentBase Info VName StructType
i Exp
uboundexp -> do
        uboundexp' <- Exp -> TermTypeM Exp
checkExp Exp
uboundexp
        it <- expType uboundexp'
        let i' = IdentBase Info VName StructType
i {identType = Info it}
        bindingIdent i' $
          -- Note: bindingParam, not bindingPat, because we must
          -- preserve the Diet of the pattern as determined by the
          -- unsized type checker (the loop parameter may be
          -- consumable).
          bindingParam looppat loop_t $ \Pat ParamType
looppat' -> 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
            loopbody' <- Exp -> TermTypeM Exp
checkExp Exp
loopbody
            (sparams, looppat'') <- checkLoopReturnSize looppat' loopbody'
            pure
              ( sparams,
                looppat'',
                For i' uboundexp',
                loopbody'
              )
      ForIn PatBase Info VName StructType
xpat Exp
e -> do
        e' <- Exp -> TermTypeM Exp
checkExp Exp
e
        t <- expTypeFully e'
        case t of
          StructType
_
            | Just StructType
t' <- Level -> StructType -> Maybe StructType
forall dim o. Level -> TypeBase dim o -> Maybe (TypeBase dim o)
peelArray Level
1 StructType
t ->
                [SizeBinder VName]
-> PatBase Info VName StructType
-> StructType
-> (Pat ParamType
    -> TermTypeM
         ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
-> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
forall o a.
[SizeBinder VName]
-> Pat (TypeBase Exp o)
-> StructType
-> (Pat ParamType -> TermTypeM a)
-> TermTypeM a
bindingPat [] PatBase Info 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' ->
                  Pat ParamType
-> StructType
-> (Pat ParamType
    -> TermTypeM
         ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
-> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
forall a.
Pat ParamType
-> StructType -> (Pat ParamType -> TermTypeM a) -> TermTypeM a
bindingParam Pat ParamType
looppat StructType
loop_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
looppat' -> 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
                    loopbody' <- Exp -> TermTypeM Exp
checkExp Exp
loopbody
                    (sparams, looppat'') <- checkLoopReturnSize looppat' loopbody'
                    pure
                      ( sparams,
                        looppat'',
                        ForIn (fmap toStruct xpat') e',
                        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 (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
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 Exp
cond ->
        Pat ParamType
-> StructType
-> (Pat ParamType
    -> TermTypeM
         ([VName], Pat ParamType, LoopFormBase Info VName, Exp))
-> TermTypeM ([VName], Pat ParamType, LoopFormBase Info VName, Exp)
forall a.
Pat ParamType
-> StructType -> (Pat ParamType -> TermTypeM a) -> TermTypeM a
bindingParam Pat ParamType
looppat StructType
loop_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
looppat' ->
          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
            cond' <-
              Exp -> TermTypeM Exp
checkExp Exp
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 NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
Bool)
            loopbody' <- checkExp loopbody
            (sparams, looppat'') <- checkLoopReturnSize looppat' loopbody'
            pure
              ( sparams,
                looppat'',
                While cond',
                loopbody'
              )

  -- dim handling (4)
  wellTypedLoopArg Initial sparams looppat' loopinit'

  (loopt, retext) <-
    freshDimsInType
      (mkUsage loc "inference of loop result type")
      (Rigid RigidLoop)
      "loop"
      sparams
      (patternType looppat')
  pure
    ( (sparams, looppat', LoopInitExplicit loopinit', form', loopbody'),
      AppRes (toStruct loopt) retext
    )

-- Note [Loop size inference]
--
-- A loop
--
--   loop p = e_init (while/for ...) do e_body
--
-- behaves like defining a function @f p = e_body@ and then calling it
-- repeatedly as @f (f ... (f e_init))@. Size-checking mirrors that: for each
-- dimension of the loop parameter @p@ we must infer whether it stays fixed
-- across iterations (*invariant*, so it can name a size that is meaningful
-- outside the loop) or may change (*variant*, so it must be existentially
-- quantified in the loop's result type). This is done in four steps, tagged
-- "dim handling (1)".."(4)" in 'checkLoop':
--
-- (1) Instantiate every size in the type of @e_init@ with a fresh nonrigid
--     variable (via 'allDimsFreshInType'), giving @loop_t@ and a map from each
--     fresh variable to the initial size it replaced. These fresh variables
--     ('new_dims') are the sizes whose variance we must determine. Distinct
--     occurrences get distinct variables even when they replace the same
--     initial size, so that e.g. two parameters that both start out @[n]@ can
--     still evolve independently.
--
-- (2) Check @e_body@ with @p@ bound at type @loop_t@, then compare the body's
--     result type against the parameter type dimension-by-dimension (with
--     'matchDims' and 'onDims'). For a new_dim @v@ that replaced initial size
--     @e'@, look at the size @d@ the body produced in that same position and
--     classify @v@ ('DimClass'):
--
--       * If @d@ is @e'@ itself, then @v@ is invariant (@Reproduces Nothing@).
--
--       * If @d@ is another new_dim @u@ that shares @v@'s initial size, then
--         @v@ is invariant iff @u@ is (@Reproduces (Just u)@). This is the
--         swap/rotation case: parameters exchanging equally-sized arrays stay
--         fixed, but if an array one of them receives has been resized, they
--         become variant.
--
--       * If @d@ is a genuinely new size, then @v@ is variant (@New@).
--
--     Variance can be mutual (in @loop (a,b) = ... in (b,a)@ each of @a@,@b@
--     copies the other's size), so we take the variant set to be the least set
--     that contains every @New@ position and is closed under the copies-from
--     dependency. Parameters that only copy from one another, with no @New@
--     seed feeding the cycle, are therefore never made variant -- a plain swap
--     of two @[n]@ arrays keeps its precise size, while resizing one of them
--     makes both existential. The variant new_dims become the loop's size
--     parameters ('sparams'); the invariant ones are substituted back to the
--     initial size they stand for.
--
-- (3) We now conceptually have a function parameter and return type.  Check,
--     as if calling that function, that the parameter type admits the body
--     result.
--
-- (4) Likewise check that it admits the initial values; this yields the loop's
--     overall (rigid) result type.