-- | Facilities for type-checking Futhark terms.  Checking a term
-- requires a little more context to track consumption and such.
--
-- Type inference is implemented through a variation of
-- Hindley-Milner.  The main complication is supporting the rich
-- number of built-in language constructs, as well as consumption
-- and freshness.  This is mostly done in an ad hoc way, and many programs
-- will require the programmer to fall back on type annotations.
--
-- The strategy is to split type checking into sveral (main) passes:
--
-- 1) A size-agnostic pass implemented in
-- "Language.Futhark.TypeChecker.Terms.Unsized".
--
-- 2) Pass (1) has given us a program where we know the types of
-- everything, but the sizes of nothing. Pass (2) then does
-- essentially size inference, with the benefit of already knowing the
-- full unsized type of everything. This is done using a syntax-driven
-- approach, similar to Algorithm W.
--
-- 3) The program is then checked for violation of the consumption
-- and freshness rules, which is implemented in
-- "Language.Futhark.TypeChecker.Consumption".
module Language.Futhark.TypeChecker.Terms
  ( checkOneExp,
    checkSizeExp,
    checkFunDef,
  )
where

import Control.Monad
import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifunctor
import Data.Bitraversable
import Data.Char (isAscii)
import Data.Either
import Data.List (delete, find, genericLength, partition)
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.Util (mapAccumLM, nubOrd, topologicalSort)
import Futhark.Util.Pretty hiding (space)
import Language.Futhark
import Language.Futhark.Primitive (intByteSize)
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Consumption qualified as Consumption
import Language.Futhark.TypeChecker.Match
import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupAbsTy, lookupMod)
import Language.Futhark.TypeChecker.Terms.Loop
import Language.Futhark.TypeChecker.Terms.Monad
import Language.Futhark.TypeChecker.Terms.Pat
import Language.Futhark.TypeChecker.Terms.Unsized qualified as Unsized
import Language.Futhark.TypeChecker.Types
import Language.Futhark.TypeChecker.Unify
import Prelude hiding (mod)

hasBinding :: Exp -> Bool
hasBinding :: Exp -> Bool
hasBinding Lambda {} = Bool
True
hasBinding (AppExp LetPat {} Info AppRes
_) = Bool
True
hasBinding (AppExp LetFun {} Info AppRes
_) = Bool
True
hasBinding (AppExp Loop {} Info AppRes
_) = Bool
True
hasBinding (AppExp LetWith {} Info AppRes
_) = Bool
True
hasBinding (AppExp Match {} Info AppRes
_) = Bool
True
hasBinding Exp
e = Maybe Exp -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe Exp -> Bool) -> Maybe Exp -> Bool
forall a b. (a -> b) -> a -> b
$ ASTMapper Maybe -> Exp -> Maybe Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper Maybe
m Exp
e
  where
    m :: ASTMapper Maybe
m =
      ASTMapper Maybe
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = \Exp
e' -> if Exp -> Bool
hasBinding Exp
e' then Maybe Exp
forall a. Maybe a
Nothing else Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e'}

--- Basic checking

-- | Determine if the two types are identical, ignoring annotations.
-- Mismatched dimensions are turned into fresh rigid type variables.
-- Causes a 'TypeError' if they fail to match, and otherwise returns
-- one of them.
unifyBranchTypes :: SrcLoc -> StructType -> StructType -> TermTypeM (StructType, [VName])
unifyBranchTypes :: SrcLoc
-> StructType -> StructType -> TermTypeM (StructType, [VName])
unifyBranchTypes SrcLoc
loc StructType
t1 StructType
t2 =
  Checking
-> TermTypeM (StructType, [VName])
-> TermTypeM (StructType, [VName])
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (StructType -> StructType -> Checking
CheckingBranches StructType
t1 StructType
t2) (TermTypeM (StructType, [VName])
 -> TermTypeM (StructType, [VName]))
-> TermTypeM (StructType, [VName])
-> TermTypeM (StructType, [VName])
forall a b. (a -> b) -> a -> b
$
    Usage
-> StructType -> StructType -> TermTypeM (StructType, [VName])
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m (StructType, [VName])
unifyMostCommon (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"unification of branch results") StructType
t1 StructType
t2

unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (StructType, [VName])
unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (StructType, [VName])
unifyBranches SrcLoc
loc Exp
e1 Exp
e2 = do
  e1_t <- Exp -> TermTypeM StructType
expType Exp
e1
  e2_t <- expType e2
  unifyBranchTypes loc e1_t e2_t

sliceShape ::
  Maybe (SrcLoc, Rigidity) ->
  [DimIndex] ->
  TypeBase Size o ->
  TermTypeM (TypeBase Size o, [VName])
sliceShape :: forall o.
Maybe (SrcLoc, Rigidity)
-> [DimIndex]
-> TypeBase Exp o
-> TermTypeM (TypeBase Exp o, [VName])
sliceShape Maybe (SrcLoc, Rigidity)
r [DimIndex]
slice t :: TypeBase Exp o
t@(Array o
o (Shape [Exp]
orig_dims) ScalarTypeBase Exp NoMode
et) =
  StateT [VName] TermTypeM (TypeBase Exp o)
-> [VName] -> TermTypeM (TypeBase Exp o, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ([Exp] -> TypeBase Exp o
setDims ([Exp] -> TypeBase Exp o)
-> StateT [VName] TermTypeM [Exp]
-> StateT [VName] TermTypeM (TypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DimIndex] -> [Exp] -> StateT [VName] TermTypeM [Exp]
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
[DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
slice [Exp]
orig_dims) []
  where
    setDims :: [Exp] -> TypeBase Exp o
setDims [] = Level -> TypeBase Exp o -> TypeBase Exp o
forall dim o. Level -> TypeBase dim o -> TypeBase dim o
stripArray ([Exp] -> Level
forall a. [a] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [Exp]
orig_dims) TypeBase Exp o
t
    setDims [Exp]
dims' = o -> Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Exp]
dims') ScalarTypeBase Exp NoMode
et

    -- If the result is supposed to be a nonrigid size variable, then
    -- don't bother trying to create non-existential sizes.  This is
    -- necessary to make programs type-check without too much
    -- ceremony; see e.g. tests/inplace5.fut.
    isRigid :: Rigidity -> Bool
isRigid Rigid {} = Bool
True
    isRigid Rigidity
_ = Bool
False
    refine_sizes :: Bool
refine_sizes = Bool
-> ((SrcLoc, Rigidity) -> Bool) -> Maybe (SrcLoc, Rigidity) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Rigidity -> Bool
isRigid (Rigidity -> Bool)
-> ((SrcLoc, Rigidity) -> Rigidity) -> (SrcLoc, Rigidity) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SrcLoc, Rigidity) -> Rigidity
forall a b. (a, b) -> b
snd) Maybe (SrcLoc, Rigidity)
r

    sliceSize :: Exp -> Maybe Exp -> Maybe Exp -> Maybe Exp -> t TermTypeM Exp
sliceSize Exp
orig_d Maybe Exp
i Maybe Exp
j Maybe Exp
stride =
      case Maybe (SrcLoc, Rigidity)
r of
        Just (SrcLoc
loc, Rigid RigidSource
_) -> do
          (d, ext) <-
            TermTypeM (Exp, Maybe VName) -> t TermTypeM (Exp, Maybe 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 (TermTypeM (Exp, Maybe VName) -> t TermTypeM (Exp, Maybe VName))
-> (SizeSource -> TermTypeM (Exp, Maybe VName))
-> SizeSource
-> t TermTypeM (Exp, Maybe VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> SizeSource -> TermTypeM (Exp, Maybe VName)
extSize SrcLoc
loc (SizeSource -> t TermTypeM (Exp, Maybe VName))
-> SizeSource -> t TermTypeM (Exp, Maybe VName)
forall a b. (a -> b) -> a -> b
$
              Maybe Exp
-> Maybe (ExpBase NoInfo VName)
-> Maybe (ExpBase NoInfo VName)
-> Maybe (ExpBase NoInfo VName)
-> SizeSource
SourceSlice Maybe Exp
orig_d' (Exp -> ExpBase NoInfo VName
bareExp (Exp -> ExpBase NoInfo VName)
-> Maybe Exp -> Maybe (ExpBase NoInfo VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp
i) (Exp -> ExpBase NoInfo VName
bareExp (Exp -> ExpBase NoInfo VName)
-> Maybe Exp -> Maybe (ExpBase NoInfo VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp
j) (Exp -> ExpBase NoInfo VName
bareExp (Exp -> ExpBase NoInfo VName)
-> Maybe Exp -> Maybe (ExpBase NoInfo VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp
stride)
          modify (maybeToList ext ++)
          pure d
        Just (SrcLoc
loc, Rigidity
Nonrigid) ->
          TermTypeM Exp -> t TermTypeM Exp
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (TermTypeM Exp -> t TermTypeM Exp)
-> TermTypeM Exp -> t TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
            (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName SrcLoc
loc (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName
              (VName -> Exp) -> TermTypeM VName -> TermTypeM Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Usage -> Name -> TermTypeM VName
forall (m :: * -> *). MonadUnify m => Usage -> Name -> m VName
newFlexibleDim (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"size of slice") Name
"slice_dim"
        Maybe (SrcLoc, Rigidity)
Nothing -> do
          v <- TermTypeM VName -> t TermTypeM 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 (TermTypeM VName -> t TermTypeM VName)
-> TermTypeM VName -> t TermTypeM VName
forall a b. (a -> b) -> a -> b
$ Name -> TermTypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID Name
"slice_anydim"
          modify (v :)
          pure $ sizeFromName (qualName v) mempty
      where
        -- The original size does not matter if the slice is fully specified.
        orig_d' :: Maybe Exp
orig_d'
          | Maybe Exp -> Bool
forall a. Maybe a -> Bool
isJust Maybe Exp
i, Maybe Exp -> Bool
forall a. Maybe a -> Bool
isJust Maybe Exp
j = Maybe Exp
forall a. Maybe a
Nothing
          | Bool
otherwise = Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
orig_d

    warnIfBinding :: Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
warnIfBinding Bool
binds Exp
d Maybe Exp
i Maybe Exp
j Maybe Exp
stride Exp
size =
      if Bool
binds
        then do
          TermTypeM () -> t TermTypeM ()
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (TermTypeM () -> t TermTypeM ())
-> (Doc () -> TermTypeM ()) -> Doc () -> t TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Doc () -> TermTypeM ()
forall loc. Located loc => loc -> Doc () -> TermTypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
size) (Doc () -> t TermTypeM ()) -> Doc () -> t TermTypeM ()
forall a b. (a -> b) -> a -> b
$
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink
              Doc ()
"size-expression-bind"
              Doc ()
"Size expression with binding is replaced by unknown size."
          (:) (Exp -> [Exp] -> [Exp])
-> t TermTypeM Exp -> t TermTypeM ([Exp] -> [Exp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Maybe Exp -> Maybe Exp -> Maybe Exp -> t TermTypeM Exp
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
Exp -> Maybe Exp -> Maybe Exp -> Maybe Exp -> t TermTypeM Exp
sliceSize Exp
d Maybe Exp
i Maybe Exp
j Maybe Exp
stride
        else ([Exp] -> [Exp]) -> t TermTypeM ([Exp] -> [Exp])
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
size Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
:)

    adjustDims :: [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims (DimFix {} : [DimIndex]
idxes') (Exp
_ : [Exp]
dims) =
      [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
idxes' [Exp]
dims
    -- Pat match some known slices to be non-existential.
    adjustDims (DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
stride : [DimIndex]
idxes') (Exp
d : [Exp]
dims)
      | Bool
refine_sizes,
        Bool -> (Exp -> Bool) -> Maybe Exp -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True ((Maybe Int64 -> Maybe Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Maybe Int64
forall a. a -> Maybe a
Just Int64
0) (Maybe Int64 -> Bool) -> (Exp -> Maybe Int64) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Maybe Int64
isInt64) Maybe Exp
i,
        Bool -> (Exp -> Bool) -> Maybe Exp -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True ((Maybe Int64 -> Maybe Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Maybe Int64
forall a. a -> Maybe a
Just Int64
1) (Maybe Int64 -> Bool) -> (Exp -> Maybe Int64) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Maybe Int64
isInt64) Maybe Exp
stride = do
          let binds :: Bool
binds = Bool -> (Exp -> Bool) -> Maybe Exp -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Exp -> Bool
hasBinding Maybe Exp
j
          Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
warnIfBinding Bool
binds Exp
d Maybe Exp
i Maybe Exp
j Maybe Exp
stride (Exp -> Maybe Exp -> Exp
forall a. a -> Maybe a -> a
fromMaybe Exp
d Maybe Exp
j)
            t TermTypeM ([Exp] -> [Exp])
-> t TermTypeM [Exp] -> t TermTypeM [Exp]
forall a b. t TermTypeM (a -> b) -> t TermTypeM a -> t TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
idxes' [Exp]
dims
    adjustDims ((DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
stride) : [DimIndex]
idxes') (Exp
d : [Exp]
dims)
      | Bool
refine_sizes,
        Just Exp
i' <- Maybe Exp
i, -- if i ~ 0, previous case
        Bool -> (Exp -> Bool) -> Maybe Exp -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True ((Maybe Int64 -> Maybe Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Maybe Int64
forall a. a -> Maybe a
Just Int64
1) (Maybe Int64 -> Bool) -> (Exp -> Maybe Int64) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Maybe Int64
isInt64) Maybe Exp
stride = do
          let j' :: Exp
j' = Exp -> Maybe Exp -> Exp
forall a. a -> Maybe a -> a
fromMaybe Exp
d Maybe Exp
j
              binds :: Bool
binds = Exp -> Bool
hasBinding Exp
j' Bool -> Bool -> Bool
|| Exp -> Bool
hasBinding Exp
i'
          Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
warnIfBinding Bool
binds Exp
d Maybe Exp
i Maybe Exp
j Maybe Exp
stride (Exp -> Exp -> Exp
sizeMinus Exp
j' Exp
i')
            t TermTypeM ([Exp] -> [Exp])
-> t TermTypeM [Exp] -> t TermTypeM [Exp]
forall a b. t TermTypeM (a -> b) -> t TermTypeM a -> t TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
idxes' [Exp]
dims
    -- stride == -1
    adjustDims ((DimSlice Maybe Exp
Nothing Maybe Exp
Nothing Maybe Exp
stride) : [DimIndex]
idxes') (Exp
d : [Exp]
dims)
      | Bool
refine_sizes,
        Bool -> (Exp -> Bool) -> Maybe Exp -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True ((Maybe Int64 -> Maybe Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Maybe Int64
forall a. a -> Maybe a
Just (-Int64
1)) (Maybe Int64 -> Bool) -> (Exp -> Maybe Int64) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Maybe Int64
isInt64) Maybe Exp
stride =
          (Exp
d Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
:) ([Exp] -> [Exp]) -> t TermTypeM [Exp] -> t TermTypeM [Exp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
idxes' [Exp]
dims
    adjustDims ((DimSlice (Just Exp
i) (Just Exp
j) Maybe Exp
stride) : [DimIndex]
idxes') (Exp
d : [Exp]
dims)
      | Bool
refine_sizes,
        Bool -> (Exp -> Bool) -> Maybe Exp -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True ((Maybe Int64 -> Maybe Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Maybe Int64
forall a. a -> Maybe a
Just (-Int64
1)) (Maybe Int64 -> Bool) -> (Exp -> Maybe Int64) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Maybe Int64
isInt64) Maybe Exp
stride = do
          let binds :: Bool
binds = Exp -> Bool
hasBinding Exp
i Bool -> Bool -> Bool
|| Exp -> Bool
hasBinding Exp
j
          Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
Bool
-> Exp
-> Maybe Exp
-> Maybe Exp
-> Maybe Exp
-> Exp
-> t TermTypeM ([Exp] -> [Exp])
warnIfBinding Bool
binds Exp
d (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
i) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
j) Maybe Exp
stride (Exp -> Exp -> Exp
sizeMinus Exp
i Exp
j)
            t TermTypeM ([Exp] -> [Exp])
-> t TermTypeM [Exp] -> t TermTypeM [Exp]
forall a b. t TermTypeM (a -> b) -> t TermTypeM a -> t TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
idxes' [Exp]
dims
    -- existential
    adjustDims ((DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
stride) : [DimIndex]
idxes') (Exp
d : [Exp]
dims) =
      (:) (Exp -> [Exp] -> [Exp])
-> t TermTypeM Exp -> t TermTypeM ([Exp] -> [Exp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Maybe Exp -> Maybe Exp -> Maybe Exp -> t TermTypeM Exp
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
Exp -> Maybe Exp -> Maybe Exp -> Maybe Exp -> t TermTypeM Exp
sliceSize Exp
d Maybe Exp
i Maybe Exp
j Maybe Exp
stride t TermTypeM ([Exp] -> [Exp])
-> t TermTypeM [Exp] -> t TermTypeM [Exp]
forall a b. t TermTypeM (a -> b) -> t TermTypeM a -> t TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [DimIndex] -> [Exp] -> t TermTypeM [Exp]
adjustDims [DimIndex]
idxes' [Exp]
dims
    adjustDims [DimIndex]
_ [Exp]
dims =
      [Exp] -> t TermTypeM [Exp]
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Exp]
dims

    sizeMinus :: Exp -> Exp -> Exp
sizeMinus Exp
j Exp
i =
      AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        ( (QualName VName, SrcLoc)
-> Info StructType
-> (Exp, Info (Maybe VName))
-> (Exp, Info (Maybe VName))
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
(QualName vn, SrcLoc)
-> f StructType
-> (ExpBase f vn, f (Maybe VName))
-> (ExpBase f vn, f (Maybe VName))
-> SrcLoc
-> AppExpBase f vn
BinOp
            (VName -> QualName VName
forall v. v -> QualName v
qualName (Name -> VName
intrinsicVar Name
"-"), SrcLoc
forall a. Monoid a => a
mempty)
            Info StructType
forall {d}. Info (TypeBase d NoMode)
sizeBinOpInfo
            (Exp
j, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
            (Exp
i, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
            SrcLoc
forall a. Monoid a => a
mempty
        )
        (Info AppRes -> Exp) -> Info AppRes -> Exp
forall a b. (a -> b) -> a -> b
$ AppRes -> Info AppRes
forall a. a -> Info a
Info
        (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
forall {dim} {o}. TypeBase dim o
i64 []
    i64 :: TypeBase dim o
i64 = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase dim o)
-> PrimType -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    sizeBinOpInfo :: Info (TypeBase d NoMode)
sizeBinOpInfo = TypeBase d NoMode -> Info (TypeBase d NoMode)
forall a. a -> Info a
Info (TypeBase d NoMode -> Info (TypeBase d NoMode))
-> TypeBase d NoMode -> Info (TypeBase d NoMode)
forall a b. (a -> b) -> a -> b
$ [TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [TypeBase d Diet
forall {dim} {o}. TypeBase dim o
i64, TypeBase d Diet
forall {dim} {o}. TypeBase dim o
i64] (RetTypeBase d Freshness -> TypeBase d NoMode)
-> RetTypeBase d Freshness -> TypeBase d NoMode
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase d Freshness -> RetTypeBase d Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase d Freshness
forall {dim} {o}. TypeBase dim o
i64
sliceShape Maybe (SrcLoc, Rigidity)
_ [DimIndex]
_ TypeBase Exp o
t = (TypeBase Exp o, [VName]) -> TermTypeM (TypeBase Exp o, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp o
t, [])

--- Main checkers

checkAscript ::
  SrcLoc ->
  TypeExp Exp VName ->
  Exp ->
  TermTypeM (TypeExp Exp VName, Exp)
checkAscript :: SrcLoc
-> TypeExp Exp VName -> Exp -> TermTypeM (TypeExp Exp VName, Exp)
checkAscript SrcLoc
loc TypeExp Exp VName
te Exp
e = do
  (te', decl_t, _) <- TypeExp Exp VName
-> TermTypeM (TypeExp Exp VName, ResType, [VName])
checkTypeExpNonrigid TypeExp Exp VName
te
  e' <- checkExp e
  e_t <- expTypeFully e'

  onFailure (CheckingAscription (toStruct decl_t) e_t) $
    unify (mkUsage loc "type ascription") (toStruct decl_t) e_t

  pure (te', e')

checkCoerce ::
  SrcLoc ->
  TypeExp Exp VName ->
  Exp ->
  TermTypeM (TypeExp Exp VName, StructType, Exp)
checkCoerce :: SrcLoc
-> TypeExp Exp VName
-> Exp
-> TermTypeM (TypeExp Exp VName, StructType, Exp)
checkCoerce SrcLoc
loc TypeExp Exp VName
te Exp
e = do
  (te', te_t, ext) <- TypeExp Exp VName
-> TermTypeM (TypeExp Exp VName, ResType, [VName])
checkTypeExpNonrigid TypeExp Exp VName
te
  e' <- checkExp e
  e_t <- expTypeFully e'

  te_t_nonrigid <- makeNonExtFresh ext $ toStruct te_t

  onFailure (CheckingAscription (toStruct te_t) e_t) $
    unify (mkUsage loc "size coercion") e_t te_t_nonrigid

  -- If the type expression had any anonymous dimensions, these will
  -- now be in 'ext'.  Those we keep nonrigid and unify with e_t.
  -- This ensures that 'x :> [1][]i32' does not make the second
  -- dimension unknown.  Use of matchDims is sensible because the
  -- structure of e_t' will be fully known due to the unification, and
  -- te_t because type expressions are complete.
  pure (te', toStruct te_t, e')
  where
    makeNonExtFresh :: t VName -> t Exp d -> TermTypeM (t Exp d)
makeNonExtFresh t VName
ext = (Exp -> TermTypeM Exp)
-> (d -> TermTypeM d) -> t Exp d -> TermTypeM (t Exp d)
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse Exp -> TermTypeM Exp
onDim d -> TermTypeM d
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      where
        onDim :: Exp -> TermTypeM Exp
onDim d :: Exp
d@(Var QualName VName
v Info StructType
_ SrcLoc
_)
          | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> t VName -> Bool
forall a. Eq a => a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
ext = Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d
        onDim Exp
d = do
          v <- Name -> TermTypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
"coerce"
          constrain v . Size Nothing $
            mkUsage
              loc
              "a size coercion where the underlying expression size cannot be determined"
          pure $ sizeFromName (qualName v) (srclocOf d)

-- Expressions witnessed by type, topologically sorted.
topWit :: TypeBase Exp o -> [Exp]
topWit :: forall o. TypeBase Exp o -> [Exp]
topWit = (Exp -> Exp -> Bool) -> [Exp] -> [Exp]
forall a. (a -> a -> Bool) -> [a] -> [a]
topologicalSort Exp -> Exp -> Bool
depends ([Exp] -> [Exp])
-> (TypeBase Exp o -> [Exp]) -> TypeBase Exp o -> [Exp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Exp o -> [Exp]
forall {fdim} {o}. TypeBase fdim o -> [fdim]
witnessedExps
  where
    witnessedExps :: TypeBase fdim o -> [fdim]
witnessedExps TypeBase fdim o
t = State [fdim] (TypeBase () o) -> [fdim] -> [fdim]
forall s a. State s a -> s -> s
execState ((Set VName -> DimPos -> fdim -> StateT [fdim] Identity ())
-> TypeBase fdim o -> State [fdim] (TypeBase () o)
forall (f :: * -> *) fdim tdim o.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim o -> f (TypeBase tdim o)
traverseDims Set VName -> DimPos -> fdim -> StateT [fdim] Identity ()
forall {a} {m :: * -> *} {p}.
MonadState [a] m =>
p -> DimPos -> a -> m ()
onDim TypeBase fdim o
t) [fdim]
forall a. Monoid a => a
mempty
      where
        onDim :: p -> DimPos -> a -> m ()
onDim p
_ DimPos
PosImmediate a
e = ([a] -> [a]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (a
e a -> [a] -> [a]
forall a. a -> [a] -> [a]
:)
        onDim p
_ DimPos
_ a
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    depends :: Exp -> Exp -> Bool
depends Exp
a Exp
b = (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Exp -> Exp -> Bool
sameExp Exp
b) ([Exp] -> Bool) -> [Exp] -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> [Exp]
subExps Exp
a

sizeFree ::
  (MonadUnify m) =>
  SrcLoc ->
  (Exp -> Maybe VName) ->
  TypeBase Size o ->
  m (TypeBase Size o, [VName])
sizeFree :: forall (m :: * -> *) o.
MonadUnify m =>
SrcLoc
-> (Exp -> Maybe VName)
-> TypeBase Exp o
-> m (TypeBase Exp o, [VName])
sizeFree SrcLoc
tloc Exp -> Maybe VName
expKiller TypeBase Exp o
orig_t = do
  ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
-> [(Exp, Exp)] -> StateT [VName] m (TypeBase Exp o)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *} {o} {b}.
(MonadState [VName] (t (t m)), MonadUnify m, MonadTrans t,
 MonadTrans t, MonadReader [(Exp, Exp)] (t (t m)), Monad (t m)) =>
TypeBase Exp o -> t (t m) b -> t (t m) b
toBeReplaced TypeBase Exp o
orig_t (ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
 -> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o))
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall (m :: * -> *) o.
MonadUnify m =>
TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
onType TypeBase Exp o
orig_t) [(Exp, Exp)]
forall a. Monoid a => a
mempty StateT [VName] m (TypeBase Exp o)
-> [VName] -> m (TypeBase Exp o, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
`runStateT` [VName]
forall a. Monoid a => a
mempty
  where
    lookReplacement :: Exp -> t (Exp, b) -> Maybe b
lookReplacement Exp
e t (Exp, b)
repl = (Exp, b) -> b
forall a b. (a, b) -> b
snd ((Exp, b) -> b) -> Maybe (Exp, b) -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, b) -> Bool) -> t (Exp, b) -> Maybe (Exp, b)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
L.find (Exp -> Exp -> Bool
sameExp Exp
e (Exp -> Bool) -> ((Exp, b) -> Exp) -> (Exp, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp, b) -> Exp
forall a b. (a, b) -> a
fst) t (Exp, b)
repl
    expReplace :: t (Exp, Exp) -> Exp -> Exp
expReplace t (Exp, Exp)
mapping Exp
e
      | Just Exp
e' <- Exp -> t (Exp, Exp) -> Maybe Exp
forall {t :: * -> *} {b}.
Foldable t =>
Exp -> t (Exp, b) -> Maybe b
lookReplacement Exp
e t (Exp, Exp)
mapping = Exp
e'
      | Bool
otherwise = Identity Exp -> Exp
forall a. Identity a -> a
runIdentity (Identity Exp -> Exp) -> Identity Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ASTMapper Identity -> Exp -> Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper Identity
mapper Exp
e
      where
        mapper :: ASTMapper Identity
mapper = ASTMapper Identity
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = pure . expReplace mapping}

    replacing :: Exp -> t (t m) Exp
replacing Exp
e = do
      e' <- (t (Exp, Exp) -> Exp) -> t (t m) Exp
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (t (Exp, Exp) -> Exp -> Exp
forall {t :: * -> *}. Foldable t => t (Exp, Exp) -> Exp -> Exp
`expReplace` Exp
e)
      case expKiller e' of
        Maybe VName
Nothing -> Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e'
        Just VName
cause -> do
          vn <- t m VName -> t (t m) VName
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (t m VName -> t (t m) VName) -> t m VName -> t (t m) VName
forall a b. (a -> b) -> a -> b
$ m VName -> t m VName
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ SrcLoc -> RigidSource -> Name -> m VName
forall a. Located a => a -> RigidSource -> Name -> m VName
forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim SrcLoc
tloc (Loc -> VName -> RigidSource
RigidOutOfScope (Exp -> Loc
forall a. Located a => a -> Loc
locOf Exp
e) VName
cause) Name
"d"
          modify (vn :)
          pure $ sizeFromName (qualName vn) (srclocOf e)

    toBeReplaced :: TypeBase Exp o -> t (t m) b -> t (t m) b
toBeReplaced TypeBase Exp o
t t (t m) b
m' = (t (t m) b -> Exp -> t (t m) b) -> t (t m) b -> [Exp] -> t (t m) b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl t (t m) b -> Exp -> t (t m) b
forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *} {b}.
(MonadState [VName] (t (t m)), MonadUnify m, MonadTrans t,
 MonadTrans t, MonadReader [(Exp, Exp)] (t (t m)), Monad (t m)) =>
t (t m) b -> Exp -> t (t m) b
f t (t m) b
m' ([Exp] -> t (t m) b) -> [Exp] -> t (t m) b
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> [Exp]
forall o. TypeBase Exp o -> [Exp]
topWit TypeBase Exp o
t
      where
        f :: t (t m) b -> Exp -> t (t m) b
f t (t m) b
m Exp
e = do
          e' <- Exp -> t (t m) Exp
forall {t :: * -> *} {t :: (* -> *) -> * -> *}
       {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadReader (t (Exp, Exp)) (t (t m)), Foldable t, MonadTrans t,
 MonadTrans t, Monad (t m), MonadUnify m,
 MonadState [VName] (t (t m))) =>
Exp -> t (t m) Exp
replacing Exp
e
          local ((e, e') :) m

    onScalar :: ScalarTypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
onScalar (Record Map Name (TypeBase Exp o)
fs) =
      Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o)
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (Map Name (TypeBase Exp o))
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp o
 -> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o))
-> Map Name (TypeBase Exp o)
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (Map Name (TypeBase Exp o))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall (m :: * -> *) o.
MonadUnify m =>
TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
onType Map Name (TypeBase Exp o)
fs
    onScalar (Sum Map Name [TypeBase Exp o]
cs) =
      Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o)
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (Map Name [TypeBase Exp o])
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TypeBase Exp o]
 -> ReaderT [(Exp, Exp)] (StateT [VName] m) [TypeBase Exp o])
-> Map Name [TypeBase Exp o]
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (Map Name [TypeBase Exp o])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (([TypeBase Exp o]
  -> ReaderT [(Exp, Exp)] (StateT [VName] m) [TypeBase Exp o])
 -> Map Name [TypeBase Exp o]
 -> ReaderT
      [(Exp, Exp)] (StateT [VName] m) (Map Name [TypeBase Exp o]))
-> ((TypeBase Exp o
     -> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o))
    -> [TypeBase Exp o]
    -> ReaderT [(Exp, Exp)] (StateT [VName] m) [TypeBase Exp o])
-> (TypeBase Exp o
    -> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o))
-> Map Name [TypeBase Exp o]
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (Map Name [TypeBase Exp o])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp o
 -> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o))
-> [TypeBase Exp o]
-> ReaderT [(Exp, Exp)] (StateT [VName] m) [TypeBase Exp o]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse) TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall (m :: * -> *) o.
MonadUnify m =>
TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
onType Map Name [TypeBase Exp o]
cs
    onScalar (Arrow o
as PName
pn Diet
d StructType
argT (RetType [VName]
dims ResType
retT)) = do
      argT' <- StructType -> ReaderT [(Exp, Exp)] (StateT [VName] m) StructType
forall (m :: * -> *) o.
MonadUnify m =>
TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
onType StructType
argT
      old_bound <- get
      retT' <- toBeReplaced retT $ onType retT
      rl <- state $ L.partition (`notElem` old_bound)
      let dims' = [VName]
dims [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
rl
      pure $ Arrow as pn d argT' (RetType dims' retT')
    onScalar (TypeVar o
o QualName VName
v [TypeArg Exp]
args) =
      o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
v ([TypeArg Exp] -> ScalarTypeBase Exp o)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) [TypeArg Exp]
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg Exp
 -> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeArg Exp))
-> [TypeArg Exp]
-> ReaderT [(Exp, Exp)] (StateT [VName] m) [TypeArg Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeArg Exp
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeArg Exp)
forall {m :: * -> *}.
MonadUnify m =>
TypeArg Exp
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeArg Exp)
onTypeArg [TypeArg Exp]
args
      where
        onTypeArg :: TypeArg Exp
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeArg Exp)
onTypeArg (TypeArgDim Exp
d) = Exp -> TypeArg Exp
forall dim. dim -> TypeArg dim
TypeArgDim (Exp -> TypeArg Exp)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) Exp
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeArg Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> ReaderT [(Exp, Exp)] (StateT [VName] m) Exp
forall {t :: * -> *} {t :: (* -> *) -> * -> *}
       {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadReader (t (Exp, Exp)) (t (t m)), Foldable t, MonadTrans t,
 MonadTrans t, Monad (t m), MonadUnify m,
 MonadState [VName] (t (t m))) =>
Exp -> t (t m) Exp
replacing Exp
d
        onTypeArg (TypeArgType StructType
ty) = StructType -> TypeArg Exp
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (StructType -> TypeArg Exp)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) StructType
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeArg Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> ReaderT [(Exp, Exp)] (StateT [VName] m) StructType
forall (m :: * -> *) o.
MonadUnify m =>
TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
onType StructType
ty
    onScalar (Prim PrimType
pt) = ScalarTypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
forall a. a -> ReaderT [(Exp, Exp)] (StateT [VName] m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ScalarTypeBase Exp o
 -> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o))
-> ScalarTypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
pt

    onType ::
      (MonadUnify m) =>
      TypeBase Size o ->
      ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Size o)
    onType :: forall (m :: * -> *) o.
MonadUnify m =>
TypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
onType (Array o
o Shape Exp
shape ScalarTypeBase Exp NoMode
scalar) =
      o -> Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o (Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (Shape Exp)
-> ReaderT
     [(Exp, Exp)]
     (StateT [VName] m)
     (ScalarTypeBase Exp NoMode -> TypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> ReaderT [(Exp, Exp)] (StateT [VName] m) Exp)
-> Shape Exp -> ReaderT [(Exp, Exp)] (StateT [VName] m) (Shape Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Shape a -> f (Shape b)
traverse Exp -> ReaderT [(Exp, Exp)] (StateT [VName] m) Exp
forall {t :: * -> *} {t :: (* -> *) -> * -> *}
       {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadReader (t (Exp, Exp)) (t (t m)), Foldable t, MonadTrans t,
 MonadTrans t, Monad (t m), MonadUnify m,
 MonadState [VName] (t (t m))) =>
Exp -> t (t m) Exp
replacing Shape Exp
shape ReaderT
  [(Exp, Exp)]
  (StateT [VName] m)
  (ScalarTypeBase Exp NoMode -> TypeBase Exp o)
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp NoMode)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall a b.
ReaderT [(Exp, Exp)] (StateT [VName] m) (a -> b)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) a
-> ReaderT [(Exp, Exp)] (StateT [VName] m) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ScalarTypeBase Exp NoMode
-> ReaderT
     [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp NoMode)
forall {m :: * -> *} {o}.
MonadUnify m =>
ScalarTypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
onScalar ScalarTypeBase Exp NoMode
scalar
    onType (Scalar ScalarTypeBase Exp o
ty) =
      ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScalarTypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
forall {m :: * -> *} {o}.
MonadUnify m =>
ScalarTypeBase Exp o
-> ReaderT [(Exp, Exp)] (StateT [VName] m) (ScalarTypeBase Exp o)
onScalar ScalarTypeBase Exp o
ty

-- Remove unknown sizes from function body types before we perform
-- let-generalisation. This is because if a function is inferred to return
-- something of type '[x+y]t' where 'x' or 'y' are unknown, we want to turn that
-- into '[z]t', where 'z' is a fresh unknown, which is then by
-- let-generalisation turned into '?[z].[z]t'.
unscopeUnknown ::
  TypeBase Size o ->
  TermTypeM (TypeBase Size o)
unscopeUnknown :: forall o. TypeBase Exp o -> TermTypeM (TypeBase Exp o)
unscopeUnknown TypeBase Exp o
t = do
  constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints
  -- The killer only ever fires on an unknown-size variable, so if none occurs
  -- free in the type there is nothing to do and we can skip the traversal (and
  -- the witness computation) entirely.
  if not (any (isUnknown constraints) (fvVars (freeInType t)))
    then pure t
    else -- These sizes will be immediately turned into existentials, so we do
    -- not need to care about their location.
      fst <$> sizeFree mempty (expKiller constraints) t
  where
    expKiller :: Map VName (a, Constraint) -> Exp -> Maybe VName
expKiller Map VName (a, Constraint)
_ Var {} = Maybe VName
forall a. Maybe a
Nothing
    expKiller Map VName (a, Constraint)
constraints Exp
e =
      Set VName -> Maybe VName
forall a. Set a -> Maybe a
S.lookupMin (Set VName -> Maybe VName) -> Set VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter (Map VName (a, Constraint) -> VName -> Bool
forall {k} {a}. Ord k => Map k (a, Constraint) -> k -> Bool
isUnknown Map VName (a, Constraint)
constraints) (Set VName -> Set VName) -> Set VName -> Set VName
forall a b. (a -> b) -> a -> b
$ (Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Set VName
witnesses) (Set VName -> Set VName) -> Set VName -> Set VName
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ Exp -> FV
freeInExp Exp
e
    isUnknown :: Map k (a, Constraint) -> k -> Bool
isUnknown Map k (a, Constraint)
constraints k
vn
      | Just UnknownSize {} <- (a, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((a, Constraint) -> Constraint)
-> Maybe (a, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> Map k (a, Constraint) -> Maybe (a, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
vn Map k (a, Constraint)
constraints = Bool
True
    isUnknown Map k (a, Constraint)
_ k
_ = Bool
False
    (Set VName
witnesses, Set VName
_) = StructType -> (Set VName, Set VName)
determineSizeWitnesses (StructType -> (Set VName, Set VName))
-> StructType -> (Set VName, Set VName)
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase Exp o
t

unscopeType ::
  SrcLoc ->
  [VName] ->
  TypeBase Size o ->
  TermTypeM (TypeBase Size o, [VName])
unscopeType :: forall o.
SrcLoc
-> [VName] -> TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName])
unscopeType SrcLoc
tloc [VName]
unscoped TypeBase Exp o
t
  -- Fast-path for common case where 't' has no free variables in unscoped.
  | Bool -> Bool
not ((VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
unscoped) (FV -> Set VName
fvVars (TypeBase Exp o -> FV
forall o. TypeBase Exp o -> FV
freeInType TypeBase Exp o
t))) = (TypeBase Exp o, [VName]) -> TermTypeM (TypeBase Exp o, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp o
t, [])
  | Bool
otherwise =
      SrcLoc
-> (Exp -> Maybe VName)
-> TypeBase Exp o
-> TermTypeM (TypeBase Exp o, [VName])
forall (m :: * -> *) o.
MonadUnify m =>
SrcLoc
-> (Exp -> Maybe VName)
-> TypeBase Exp o
-> m (TypeBase Exp o, [VName])
sizeFree SrcLoc
tloc ((VName -> Bool) -> Set VName -> Maybe VName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
unscoped) (Set VName -> Maybe VName)
-> (Exp -> Set VName) -> Exp -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Set VName
fvVars (FV -> Set VName) -> (Exp -> FV) -> Exp -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> FV
freeInExp) TypeBase Exp o
t

checkExp :: Exp -> TermTypeM Exp
checkExp :: Exp -> TermTypeM Exp
checkExp (Var QualName VName
qn (Info StructType
t) SrcLoc
loc) = do
  t' <- SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
qn StructType
t
  pure $ Var qn (Info t') loc
checkExp (Literal PrimValue
val SrcLoc
loc) =
  Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal PrimValue
val SrcLoc
loc
checkExp (Hole (Info StructType
t) SrcLoc
loc) = do
  t' <- SrcLoc -> StructType -> TermTypeM StructType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVarsAbsorbable SrcLoc
loc StructType
t
  pure $ Hole (Info t') loc
checkExp (StringLit [Word8]
vs SrcLoc
loc) =
  Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ [Word8] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [Word8] -> SrcLoc -> ExpBase f vn
StringLit [Word8]
vs SrcLoc
loc
checkExp (IntLit Integer
val (Info StructType
t) SrcLoc
loc) = do
  t' <- SrcLoc -> StructType -> TermTypeM StructType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVars SrcLoc
loc StructType
t
  pure $ IntLit val (Info t') loc
checkExp (FloatLit Double
val (Info StructType
t) SrcLoc
loc) = do
  t' <- SrcLoc -> StructType -> TermTypeM StructType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVars SrcLoc
loc StructType
t
  pure $ FloatLit val (Info t') loc
checkExp (TupLit [Exp]
es SrcLoc
loc) =
  [Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Exp] -> SrcLoc -> Exp)
-> TermTypeM [Exp] -> TermTypeM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> TermTypeM Exp) -> [Exp] -> TermTypeM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> TermTypeM Exp
checkExp [Exp]
es TermTypeM (SrcLoc -> Exp) -> TermTypeM SrcLoc -> TermTypeM Exp
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermTypeM SrcLoc
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) =
  [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> TermTypeM [FieldBase Info VName] -> TermTypeM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> TermTypeM (FieldBase Info VName))
-> [FieldBase Info VName] -> TermTypeM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM FieldBase Info VName -> TermTypeM (FieldBase Info VName)
checkField [FieldBase Info VName]
fs TermTypeM (SrcLoc -> Exp) -> TermTypeM SrcLoc -> TermTypeM Exp
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermTypeM SrcLoc
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  where
    checkField :: FieldBase Info VName -> TermTypeM (FieldBase Info VName)
checkField (RecordFieldExplicit L Name
f Exp
e SrcLoc
rloc) =
      L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit L Name
f (Exp -> SrcLoc -> FieldBase Info VName)
-> TermTypeM Exp -> TermTypeM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> TermTypeM Exp
checkExp Exp
e TermTypeM (SrcLoc -> FieldBase Info VName)
-> TermTypeM SrcLoc -> TermTypeM (FieldBase Info VName)
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermTypeM SrcLoc
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
rloc
    checkField (RecordFieldImplicit L VName
name (Info StructType
t) SrcLoc
rloc) = do
      t' <- SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar SrcLoc
rloc (VName -> QualName VName
forall v. v -> QualName v
qualName (L VName -> VName
forall a. L a -> a
unLoc L VName
name)) StructType
t
      pure $ RecordFieldImplicit name (Info t') rloc
-- No need to type check this, as these are only produced by the
-- parser if the elements are monomorphic and all match.
checkExp (ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ [PrimValue] -> PrimType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc
checkExp (ArrayLit [Exp]
all_es (Info StructType
t) SrcLoc
loc) =
  -- We only consult the type inferred by the unsized type checker
  -- for empty arrays; otherwise we use the type of the first
  -- element.  This significantly cuts down on the number of
  -- inferred types we have to instantiate for pathologically large
  -- multidimensional array literals.
  case [Exp]
all_es of
    [] -> do
      t' <- SrcLoc -> StructType -> TermTypeM StructType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVars SrcLoc
loc StructType
t
      case peelArray 1 t' of
        Just StructType
et -> do
          let t'' :: StructType
t'' = Shape Exp -> StructType -> StructType
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Exp
sizeFromInteger Integer
0 SrcLoc
forall a. Monoid a => a
mempty]) StructType
et
          Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"empty array literal") StructType
t'' StructType
t'
          Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
ArrayLit [] (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t'') SrcLoc
loc
        Maybe StructType
Nothing -> [Char] -> TermTypeM Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM Exp) -> [Char] -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"checkExp ArrayLit: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t'
    Exp
e : [Exp]
es -> do
      e' <- Exp -> TermTypeM Exp
checkExp Exp
e
      et <- expType e'
      es' <- mapM (unifies "type of first array element" et <=< checkExp) es
      let arr_t = Shape Exp -> StructType -> StructType
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Exp
sizeFromInteger ([Exp] -> Integer
forall i a. Num i => [a] -> i
genericLength [Exp]
all_es) SrcLoc
forall a. Monoid a => a
mempty]) StructType
et
      pure $ ArrayLit (e' : es') (Info arr_t) loc
checkExp (AppExp (Range Exp
start Maybe Exp
maybe_step Inclusiveness Exp
end SrcLoc
loc) Info AppRes
_) = do
  start' <- Exp -> TermTypeM Exp
checkExp Exp
start
  start_t <- expType start'
  maybe_step' <- case maybe_step of
    Maybe Exp
Nothing -> Maybe Exp -> TermTypeM (Maybe Exp)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Exp
forall a. Maybe a
Nothing
    Just Exp
step -> do
      let warning :: TermTypeM ()
warning = SrcLoc -> Doc () -> TermTypeM ()
forall loc. Located loc => loc -> Doc () -> TermTypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn SrcLoc
loc Doc ()
"First and second element of range are identical, this will produce an empty array."
      case (Exp
start, Exp
step) of
        (Literal PrimValue
x SrcLoc
_, Literal PrimValue
y SrcLoc
_) -> Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (PrimValue
x PrimValue -> PrimValue -> Bool
forall a. Eq a => a -> a -> Bool
== PrimValue
y) TermTypeM ()
warning
        (Var QualName VName
x_name Info StructType
_ SrcLoc
_, Var QualName VName
y_name Info StructType
_ SrcLoc
_) -> Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (QualName VName
x_name QualName VName -> QualName VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName
y_name) TermTypeM ()
warning
        (Exp, Exp)
_ -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> TermTypeM Exp -> TermTypeM (Maybe Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> StructType -> Exp -> TermTypeM Exp
unifies Text
"use in range expression" StructType
start_t (Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM Exp
checkExp Exp
step)

  let unifyRange Exp
e = Text -> StructType -> Exp -> TermTypeM Exp
unifies Text
"use in range expression" StructType
start_t (Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM Exp
checkExp Exp
e
  end' <- traverse unifyRange end

  end_t <- case end' of
    DownToExclusive Exp
e -> Exp -> TermTypeM StructType
expType Exp
e
    ToInclusive Exp
e -> Exp -> TermTypeM StructType
expType Exp
e
    UpToExclusive Exp
e -> Exp -> TermTypeM StructType
expType Exp
e

  -- Special case some ranges to give them a known size.
  let warnIfBinding Bool
binds Exp
size =
        if Bool
binds
          then do
            SrcLoc -> Doc () -> m ()
forall loc. Located loc => loc -> Doc () -> m ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
size) (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink
                Doc ()
"size-expression-bind"
                Doc ()
"Size expression with binding is replaced by unknown size."
            d <- SrcLoc -> RigidSource -> Name -> m VName
forall a. Located a => a -> RigidSource -> Name -> m VName
forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim SrcLoc
loc RigidSource
RigidRange Name
"range_dim"
            pure (sizeFromName (qualName d) mempty, Just d)
          else (Exp, Maybe VName) -> m (Exp, Maybe VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
size, Maybe VName
forall a. Maybe a
Nothing)
  (dim, retext) <-
    case (isInt64 start', isInt64 <$> maybe_step', end') of
      (Just Int64
0, Just (Just Int64
1), UpToExclusive Exp
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Exp -> TermTypeM (Exp, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Exp -> m (Exp, Maybe VName)
warnIfBinding (Exp -> Bool
hasBinding Exp
end'') Exp
end''
      (Just Int64
0, Maybe (Maybe Int64)
Nothing, UpToExclusive Exp
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Exp -> TermTypeM (Exp, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Exp -> m (Exp, Maybe VName)
warnIfBinding (Exp -> Bool
hasBinding Exp
end'') Exp
end''
      (Maybe Int64
_, Maybe (Maybe Int64)
Nothing, UpToExclusive Exp
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Exp -> TermTypeM (Exp, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Exp -> m (Exp, Maybe VName)
warnIfBinding (Exp -> Bool
hasBinding Exp
end'' Bool -> Bool -> Bool
|| Exp -> Bool
hasBinding Exp
start') (Exp -> TermTypeM (Exp, Maybe VName))
-> Exp -> TermTypeM (Exp, Maybe VName)
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
sizeMinus Exp
end'' Exp
start'
      (Maybe Int64
_, Maybe (Maybe Int64)
Nothing, ToInclusive Exp
end'')
        -- No stride means we assume a stride of one.
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Exp -> TermTypeM (Exp, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Exp -> m (Exp, Maybe VName)
warnIfBinding (Exp -> Bool
hasBinding Exp
end'' Bool -> Bool -> Bool
|| Exp -> Bool
hasBinding Exp
start') (Exp -> TermTypeM (Exp, Maybe VName))
-> Exp -> TermTypeM (Exp, Maybe VName)
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
sizeMinusInc Exp
end'' Exp
start'
      (Just Int64
1, Just (Just Int64
2), ToInclusive Exp
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Exp -> TermTypeM (Exp, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Exp -> m (Exp, Maybe VName)
warnIfBinding (Exp -> Bool
hasBinding Exp
end'') Exp
end''
      (Maybe Int64, Maybe (Maybe Int64), Inclusiveness Exp)
_ -> do
        d <- SrcLoc -> RigidSource -> Name -> TermTypeM VName
forall a. Located a => a -> RigidSource -> Name -> TermTypeM VName
forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim SrcLoc
loc RigidSource
RigidRange Name
"range_dim"
        pure (sizeFromName (qualName d) mempty, Just d)

  let t = Shape Exp -> StructType -> StructType
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Exp
dim]) StructType
start_t
      res = StructType -> [VName] -> AppRes
AppRes StructType
t (Maybe VName -> [VName]
forall a. Maybe a -> [a]
maybeToList Maybe VName
retext)

  pure $ AppExp (Range start' maybe_step' end' loc) (Info res)
  where
    i64 :: TypeBase dim o
i64 = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase dim o)
-> PrimType -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    mkBinOp :: Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
op StructType
t Exp
x Exp
y =
      AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        ( (QualName VName, SrcLoc)
-> Info StructType
-> (Exp, Info (Maybe VName))
-> (Exp, Info (Maybe VName))
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
(QualName vn, SrcLoc)
-> f StructType
-> (ExpBase f vn, f (Maybe VName))
-> (ExpBase f vn, f (Maybe VName))
-> SrcLoc
-> AppExpBase f vn
BinOp
            (VName -> QualName VName
forall v. v -> QualName v
qualName (Name -> VName
intrinsicVar Name
op), SrcLoc
forall a. Monoid a => a
mempty)
            Info StructType
forall {d}. Info (TypeBase d NoMode)
sizeBinOpInfo
            (Exp
x, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
            (Exp
y, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
            SrcLoc
forall a. Monoid a => a
mempty
        )
        (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
t [])
    mkSub :: Exp -> Exp -> Exp
mkSub = Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
"-" StructType
forall {dim} {o}. TypeBase dim o
i64
    mkAdd :: Exp -> Exp -> Exp
mkAdd = Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
"+" StructType
forall {dim} {o}. TypeBase dim o
i64
    sizeMinus :: Exp -> Exp -> Exp
sizeMinus Exp
j Exp
i = Exp
j Exp -> Exp -> Exp
`mkSub` Exp
i
    sizeMinusInc :: Exp -> Exp -> Exp
sizeMinusInc Exp
j Exp
i = (Exp
j Exp -> Exp -> Exp
`mkSub` Exp
i) Exp -> Exp -> Exp
`mkAdd` Integer -> SrcLoc -> Exp
sizeFromInteger Integer
1 SrcLoc
forall a. Monoid a => a
mempty
    sizeBinOpInfo :: Info (TypeBase d NoMode)
sizeBinOpInfo = TypeBase d NoMode -> Info (TypeBase d NoMode)
forall a. a -> Info a
Info (TypeBase d NoMode -> Info (TypeBase d NoMode))
-> TypeBase d NoMode -> Info (TypeBase d NoMode)
forall a b. (a -> b) -> a -> b
$ [TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [TypeBase d Diet
forall {dim} {o}. TypeBase dim o
i64, TypeBase d Diet
forall {dim} {o}. TypeBase dim o
i64] (RetTypeBase d Freshness -> TypeBase d NoMode)
-> RetTypeBase d Freshness -> TypeBase d NoMode
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase d Freshness -> RetTypeBase d Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase d Freshness
forall {dim} {o}. TypeBase dim o
i64
checkExp (Ascript Exp
e TypeExp Exp VName
te SrcLoc
loc) = do
  (te', e') <- SrcLoc
-> TypeExp Exp VName -> Exp -> TermTypeM (TypeExp Exp VName, Exp)
checkAscript SrcLoc
loc TypeExp Exp VName
te Exp
e
  pure $ Ascript e' te' loc
checkExp (Coerce Exp
e TypeExp Exp VName
te Info StructType
_ SrcLoc
loc) = do
  (te', te_t, e') <- SrcLoc
-> TypeExp Exp VName
-> Exp
-> TermTypeM (TypeExp Exp VName, StructType, Exp)
checkCoerce SrcLoc
loc TypeExp Exp VName
te Exp
e
  t <- expTypeFully e'
  t' <- matchDims (const . const pure) t te_t
  pure $ Coerce e' te' (Info t') loc
checkExp (AppExp (Apply Exp
fe NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
loc) Info AppRes
_) = do
  fe' <- Exp -> TermTypeM Exp
checkExp Exp
fe
  args' <- mapM (checkExp . snd) args
  t <- expType fe'
  let fname =
        case Exp
fe' of
          Var QualName VName
v Info StructType
_ SrcLoc
_ -> QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
v
          Exp
_ -> Maybe (QualName VName)
forall a. Maybe a
Nothing
  ((_, exts, rt), args'') <- mapAccumLM (onArg fname) (0, [], t) args'

  pure $ AppExp (Apply fe' args'' loc) $ Info $ AppRes rt exts
  where
    onArg :: Maybe (QualName VName)
-> (Level, [VName], StructType)
-> Exp
-> TermTypeM
     ((Level, [VName], StructType), (Info (Maybe VName), Exp))
onArg Maybe (QualName VName)
fname (Level
i, [VName]
all_exts, StructType
t) Exp
arg' = do
      (_, rt, argext, exts) <- SrcLoc
-> ApplyOp
-> StructType
-> Exp
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (Maybe (QualName VName)
fname, Level
i) StructType
t Exp
arg'
      pure
        ( (i + 1, all_exts <> exts, rt),
          (Info argext, arg')
        )
checkExp (AppExp (BinOp (QualName VName
op, SrcLoc
oploc) (Info StructType
op_t) (Exp
e1, Info (Maybe VName)
_) (Exp
e2, Info (Maybe VName)
_) SrcLoc
loc) Info AppRes
_) = do
  ftype <- SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar SrcLoc
oploc QualName VName
op StructType
op_t
  e1' <- checkExp e1
  e2' <- checkExp e2
  -- Note that the application to the first operand cannot fix any
  -- existential sizes, because it must by necessity be a function.
  (_, rt, p1_ext, _) <- checkApply loc (Just op, 0) ftype e1'
  (_, rt', p2_ext, retext) <- checkApply loc (Just op, 1) rt e2'

  pure $
    AppExp
      ( BinOp
          (op, oploc)
          (Info ftype)
          (e1', Info p1_ext)
          (e2', Info p2_ext)
          loc
      )
      (Info (AppRes rt' retext))
checkExp (Project Name
k Exp
e Info StructType
_ SrcLoc
loc) = do
  e' <- Exp -> TermTypeM Exp
checkExp Exp
e
  t <- expType e'
  case t of
    Scalar (Record Map Name StructType
fs)
      | Just StructType
kt <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
k Map Name StructType
fs ->
          Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
k Exp
e' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
kt) SrcLoc
loc
    StructType
_ -> [Char] -> TermTypeM Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM Exp) -> [Char] -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"checkExp Project: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Show a => a -> [Char]
show StructType
t
checkExp (Parens Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp)
-> TermTypeM Exp -> TermTypeM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> TermTypeM Exp
checkExp Exp
e TermTypeM (SrcLoc -> Exp) -> TermTypeM SrcLoc -> TermTypeM Exp
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermTypeM SrcLoc
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkExp (QualParens (QualName VName
modname, SrcLoc
modnameloc) Exp
e SrcLoc
loc) = do
  mod <- QualName VName -> TermTypeM Mod
lookupMod QualName VName
modname
  case mod of
    ModEnv Env
env -> (TermEnv -> TermEnv) -> TermTypeM Exp -> TermTypeM Exp
forall a. (TermEnv -> TermEnv) -> TermTypeM a -> TermTypeM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (TermEnv -> Env -> TermEnv
`withEnv` Env
env) (TermTypeM Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ do
      e' <- Exp -> TermTypeM Exp
checkExp Exp
e
      pure $ QualParens (modname, modnameloc) e' loc
    ModFun {} ->
      SrcLoc -> Notes -> Doc () -> TermTypeM 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 SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM Exp)
-> (Doc () -> Doc ()) -> Doc () -> TermTypeM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"module-is-parametric" (Doc () -> TermTypeM Exp) -> Doc () -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Module" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> QualName VName -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName VName -> Doc ann
pretty QualName VName
modname Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
" is a parametric module."
checkExp (Negate Exp
arg SrcLoc
loc) = do
  arg' <- Exp -> TermTypeM Exp
checkExp Exp
arg
  pure $ Negate arg' loc
checkExp (Not Exp
arg SrcLoc
loc) = do
  arg' <- Exp -> TermTypeM Exp
checkExp Exp
arg
  pure $ Not arg' loc
checkExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
loc) Info AppRes
_) = do
  e' <- Exp -> TermTypeM Exp
checkExp Exp
e

  -- Not technically an ascription, but we want the pattern to have
  -- exactly the type of 'e'.
  t <- expType e'
  bindingSizes sizes . incLevel . bindingPat sizes pat t $ \Pat ParamType
pat' -> do
    body' <- TermTypeM Exp -> TermTypeM Exp
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM Exp -> TermTypeM Exp) -> TermTypeM Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> TermTypeM Exp
checkExp Exp
body
    body_t <- expTypeFully body'

    -- If the bound expression is of type i64, then we replace the
    -- pattern name with the expression in the type of the body.
    -- Otherwise, we need to come up with unknown sizes for the
    -- sizes going out of scope.
    (body_t', retext) <-
      case (t, patNames pat') of
        (Scalar (Prim (Signed IntType
Int64)), [VName
v])
          | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> Bool
hasBinding Exp
e' -> do
              let f :: VName -> Maybe (Subst t)
f VName
x = if VName
x VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v then Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst Exp
e') else Maybe (Subst t)
forall a. Maybe a
Nothing
              (StructType, [VName]) -> TermTypeM (StructType, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
forall {t}. VName -> Maybe (Subst t)
f StructType
body_t, [])
        (StructType, [VName])
_ ->
          SrcLoc -> [VName] -> StructType -> TermTypeM (StructType, [VName])
forall o.
SrcLoc
-> [VName] -> TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName])
unscopeType SrcLoc
loc ((SizeBinder VName -> VName) -> [SizeBinder VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map SizeBinder VName -> VName
forall vn. SizeBinder vn -> vn
sizeName [SizeBinder VName]
sizes [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames Pat ParamType
pat') StructType
body_t

    pure $
      AppExp
        (LetPat sizes (fmap toStruct pat') e' body' loc)
        (Info $ AppRes body_t' retext)
checkExp (AppExp (LetFun (VName, SrcLoc)
name ([TypeParamBase VName]
tparams, [Pat ParamType]
params, Maybe (TypeExp Exp VName)
maybe_retdecl, Info ResRetType
_, Exp
e) Exp
body SrcLoc
loc) Info AppRes
_) = do
  (tparams', params', maybe_retdecl', rettype, e') <-
    (VName, Maybe (TypeExp Exp VName), [TypeParamBase VName],
 [Pat ParamType], Exp, SrcLoc)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
checkBinding ((VName, SrcLoc) -> VName
forall a b. (a, b) -> a
fst (VName, SrcLoc)
name, Maybe (TypeExp Exp VName)
maybe_retdecl, [TypeParamBase VName]
tparams, [Pat ParamType]
params, Exp
e, SrcLoc
loc)

  let entry = [TypeParamBase VName] -> StructType -> ValBinding Exp
forall dim.
[TypeParamBase VName] -> TypeBase dim NoMode -> ValBinding dim
BoundV [TypeParamBase VName]
tparams' (StructType -> ValBinding Exp) -> StructType -> ValBinding Exp
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> ResRetType -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
params' ResRetType
rettype
      bindF TermScope Exp
scope =
        TermScope Exp
scope
          { scopeVtable = M.insert (fst name) entry $ scopeVtable scope
          }
  body' <- localScope bindF $ checkExp body

  (body_t, ext) <- unscopeType loc [fst name] =<< expTypeFully body'

  pure $
    AppExp
      ( LetFun
          name
          (tparams', params', maybe_retdecl', Info rettype, e')
          body'
          loc
      )
      (Info $ AppRes body_t ext)
checkExp (AppExp (LetWith IdentBase Info VName StructType
dest IdentBase Info VName StructType
src [UpdateStep Info VName]
steps Exp
ve Exp
body SrcLoc
loc) Info AppRes
_) = do
  -- The type recorded in the AST is the unsized type from the
  -- unsized type checker; we must consult the scope to get the
  -- actual type of the source variable.
  src_t <-
    StructType -> TermTypeM StructType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully
      (StructType -> TermTypeM StructType)
-> TermTypeM StructType -> TermTypeM StructType
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar (IdentBase Info VName StructType -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf IdentBase Info VName StructType
src) (VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase Info VName StructType
src) (Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> StructType) -> Info StructType -> StructType
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> Info StructType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType IdentBase Info VName StructType
src)
  let src' = IdentBase Info VName StructType
src {identType = Info src_t}

  case mapAndUnzipM isField steps of
    Just ([UpdateStep Info VName]
steps', [Name]
names) -> do
      ve' <- Exp -> TermTypeM Exp
checkExp Exp
ve
      ve_t <- expType ve'
      updated_t <- updateFieldPath src names ve_t src_t

      let dest' = IdentBase Info VName StructType
dest {identType = Info updated_t}
      bindingIdent dest' $ do
        body' <- checkExp body
        (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'
        pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t ext)
    Maybe ([UpdateStep Info VName], [Name])
Nothing -> do
      (steps', target_t) <- SrcLoc
-> StructType
-> [UpdateStep Info VName]
-> TermTypeM ([UpdateStep Info VName], StructType)
checkUpdateSteps SrcLoc
loc StructType
src_t [UpdateStep Info VName]
steps
      ve' <- unifies "type of update target" target_t =<< checkExp ve

      let dest' = IdentBase Info VName StructType
dest {identType = Info src_t}
      bindingIdent dest' $ do
        body' <- checkExp body
        (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'
        pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t ext)
  where
    isField :: UpdateStep f vn -> Maybe (UpdateStep f vn, Name)
isField (UpdateStepField Name
f) = (UpdateStep f vn, Name) -> Maybe (UpdateStep f vn, Name)
forall a. a -> Maybe a
Just (Name -> UpdateStep f vn
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f, Name
f)
    isField UpdateStep f vn
_ = Maybe (UpdateStep f vn, Name)
forall a. Maybe a
Nothing

-- Record updates are a bit hacky, because we do not have row typing
-- (yet?).  For now, we only permit record updates where we know the
-- full type up to the field we are updating.
checkExp (Update Exp
src [UpdateStep Info VName]
steps Exp
ve Info StructType
_ SrcLoc
loc) = do
  src' <- Exp -> TermTypeM Exp
checkExp Exp
src
  src_t <- expTypeFully src'
  case mapAndUnzipM isField steps of
    Just ([UpdateStep Info VName]
steps', [Name]
names) -> do
      ve' <- Exp -> TermTypeM Exp
checkExp Exp
ve
      ve_t <- expType ve'
      updated_t <- updateFieldPath src names ve_t src_t
      pure $ Update src' steps' ve' (Info updated_t) loc
    Maybe ([UpdateStep Info VName], [Name])
Nothing -> do
      (steps', target_t) <- SrcLoc
-> StructType
-> [UpdateStep Info VName]
-> TermTypeM ([UpdateStep Info VName], StructType)
checkUpdateSteps SrcLoc
loc StructType
src_t [UpdateStep Info VName]
steps
      ve' <- unifies "type of update target" target_t =<< checkExp ve
      src_t' <- expTypeFully src'
      pure $ Update src' steps' ve' (Info src_t') loc
  where
    isField :: UpdateStep f vn -> Maybe (UpdateStep f vn, Name)
isField (UpdateStepField Name
f) = (UpdateStep f vn, Name) -> Maybe (UpdateStep f vn, Name)
forall a. a -> Maybe a
Just (Name -> UpdateStep f vn
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f, Name
f)
    isField UpdateStep f vn
_ = Maybe (UpdateStep f vn, Name)
forall a. Maybe a
Nothing
checkExp (AppExp (Index Exp
e [DimIndex]
slice SrcLoc
loc) Info AppRes
_) = do
  slice' <- [DimIndex] -> TermTypeM [DimIndex]
checkSlice [DimIndex]
slice
  e' <- checkExp e
  -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value.
  (t', retext) <-
    sliceShape (Just (loc, Rigid (RigidSlice Nothing ""))) slice'
      =<< expTypeFully e'

  pure $ AppExp (Index e' slice' loc) (Info $ AppRes t' retext)
checkExp (Assert Exp
e1 Exp
e2 Info Text
_ SrcLoc
loc) = do
  e1' <- Exp -> TermTypeM Exp
checkExp Exp
e1
  e2' <- checkExp e2
  pure $ Assert e1' e2' (Info (prettyText e1)) loc
checkExp (Lambda [Pat ParamType]
params Exp
body Maybe (TypeExp Exp VName)
rettype_te (Info (RetType [VName]
_ ResType
rt)) SrcLoc
loc) = do
  (params', body', rettype', RetType dims ty) <-
    TermTypeM
  ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType)
-> TermTypeM
     ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType)
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM
   ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType)
 -> TermTypeM
      ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
-> (([Pat ParamType]
     -> TermTypeM
          ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
    -> TermTypeM
         ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
-> ([Pat ParamType]
    -> TermTypeM
         ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
-> TermTypeM
     ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeParamBase VName]
-> [Pat ParamType]
-> ([Pat ParamType]
    -> TermTypeM
         ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
-> TermTypeM
     ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType)
forall a.
[TypeParamBase VName]
-> [Pat ParamType]
-> ([Pat ParamType] -> TermTypeM a)
-> TermTypeM a
bindingParams [] [Pat ParamType]
params (([Pat ParamType]
  -> TermTypeM
       ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
 -> TermTypeM
      ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
-> ([Pat ParamType]
    -> TermTypeM
         ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType))
-> TermTypeM
     ([Pat ParamType], Exp, Maybe (TypeExp Exp VName), ResRetType)
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
      -- The sizes of the return type are absorbable, as a lambda
      -- returns whatever type the context requires. See Note [Size
      -- Inference].
      rt' <- SrcLoc -> ResType -> TermTypeM ResType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVarsAbsorbable SrcLoc
loc ResType
rt
      rettype_checked <- traverse checkTypeExpNonrigid rettype_te
      declared_rettype <-
        case rettype_checked of
          Just (TypeExp Exp VName
_, ResType
st, [VName]
_) -> do
            Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (Exp -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Exp
body Text
"lambda return type ascription") (ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
rt') (ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
st)
            Maybe ResType -> TermTypeM (Maybe ResType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe ResType -> TermTypeM (Maybe ResType))
-> Maybe ResType -> TermTypeM (Maybe ResType)
forall a b. (a -> b) -> a -> b
$ ResType -> Maybe ResType
forall a. a -> Maybe a
Just ResType
st
          Maybe (TypeExp Exp VName, ResType, [VName])
Nothing -> Maybe ResType -> TermTypeM (Maybe ResType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ResType
forall a. Maybe a
Nothing
      (body', body_t) <- checkFunBody params' body declared_rettype loc

      unify (mkUsage body "inferred return type") (toStruct rt') body_t

      params'' <- mapM updateTypes params'

      -- A lambda has no let-generalisation to decide where its
      -- existential sizes go, so we infer them here - also for a
      -- declared return type, whose quantified sizes may by now have
      -- been solved to unknown sizes.
      rettype_st <-
        inferReturnSizes params'' =<< case rettype_checked of
          Just (TypeExp Exp VName
_, ResType
ret, [VName]
_) -> ResType -> TermTypeM ResType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully ResType
ret
          Maybe (TypeExp Exp VName, ResType, [VName])
Nothing -> ResType -> TermTypeM ResType
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResType -> TermTypeM ResType) -> ResType -> TermTypeM ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
body_t

      pure (params'', body', (\(TypeExp Exp VName
te, ResType
_, [VName]
_) -> TypeExp Exp VName
te) <$> rettype_checked, rettype_st)

  verifyFunctionParams Nothing params'

  (ty', dims') <- unscopeType loc dims ty

  pure $ Lambda params' body' rettype' (Info (RetType dims' ty')) loc
  where
    -- Inferring the sizes of the return type of a lambda is a lot
    -- like let-generalisation.  We wish to remove any rigid sizes
    -- that were created when checking the body, except for those that
    -- are visible in types that existed before we entered the body,
    -- are parameters, or are used in parameters.
    inferReturnSizes :: [Pat ParamType] -> ResType -> m ResRetType
inferReturnSizes [Pat ParamType]
params' ResType
ret = do
      cur_lvl <- m Level
forall (m :: * -> *). MonadUnify m => m Level
curLevel
      let named (Named VName
x, b
_, c
_) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
x
          named (PName
Unnamed, b
_, c
_) = Maybe VName
forall a. Maybe a
Nothing
          param_names = (Pat ParamType -> Maybe VName) -> [Pat ParamType] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((PName, Diet, StructType) -> Maybe VName
forall {b} {c}. (PName, b, c) -> Maybe VName
named ((PName, Diet, StructType) -> Maybe VName)
-> (Pat ParamType -> (PName, Diet, StructType))
-> Pat ParamType
-> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> (PName, Diet, StructType)
forall d. Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam) [Pat ParamType]
params'
          pos_sizes =
            StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
sizeNamesPos (StructType -> Set VName) -> StructType -> Set VName
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> ResRetType -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
params' (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] ResType
ret
          -- Only rigid sizes computed by the body can be hidden. A
          -- size that is still flexible has not been determined yet,
          -- and hiding it would sever its connection to whatever the
          -- enclosing context determines it to be.
          rigid UnknownSize {} = Bool
True
          rigid Constraint
_ = Bool
False
          hide VName
k (Level
lvl, Constraint
c) =
            Constraint -> Bool
rigid Constraint
c Bool -> Bool -> Bool
&& Level
lvl Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
>= Level
cur_lvl Bool -> Bool -> Bool
&& VName
k VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
param_names Bool -> Bool -> Bool
&& VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
pos_sizes

      hidden_sizes <-
        S.fromList . M.keys . M.filterWithKey hide <$> getConstraints

      let onDim VName
name
            | VName
name VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
hidden_sizes = VName -> Set VName
forall a. a -> Set a
S.singleton VName
name
          onDim VName
_ = Set VName
forall a. Monoid a => a
mempty

      pure $ RetType (S.toList $ foldMap onDim $ fvVars $ freeInType ret) ret
checkExp (OpSection QualName VName
op (Info StructType
op_t) SrcLoc
loc) = do
  ftype <- SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
op StructType
op_t
  pure $ OpSection op (Info ftype) loc
checkExp (OpSectionLeft QualName VName
op (Info StructType
op_t) Exp
e (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
_ (Info ResRetType, Info [VName])
_ SrcLoc
loc) = do
  ftype <- SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
op StructType
op_t
  e' <- checkExp e
  (t1, rt, argext, retext) <- checkApply loc (Just op, 0) ftype e'
  case (ftype, rt) of
    (Scalar (Arrow NoMode
_ PName
m1 Diet
d1 StructType
_ ResRetType
_), Scalar (Arrow NoMode
_ PName
m2 Diet
d2 StructType
t2 (RetType [VName]
ds ResType
rt2))) ->
      Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
        QualName VName
-> Info StructType
-> Exp
-> (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
-> (Info ResRetType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
QualName vn
-> f StructType
-> ExpBase f vn
-> (f (PName, ParamType, Maybe VName), f (PName, ParamType))
-> (f ResRetType, f [VName])
-> SrcLoc
-> ExpBase f vn
OpSectionLeft
          QualName VName
op
          (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ftype)
          Exp
e'
          ((PName, ParamType, Maybe VName)
-> Info (PName, ParamType, Maybe VName)
forall a. a -> Info a
Info (PName
m1, Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
d1 StructType
t1, Maybe VName
argext), (PName, ParamType) -> Info (PName, ParamType)
forall a. a -> Info a
Info (PName
m2, Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
d2 StructType
t2))
          (ResRetType -> Info ResRetType
forall a. a -> Info a
Info (ResRetType -> Info ResRetType) -> ResRetType -> Info ResRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ds ResType
rt2, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
retext)
          SrcLoc
loc
    (StructType, StructType)
_ ->
      SrcLoc -> Notes -> Doc () -> TermTypeM 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 SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM Exp) -> Doc () -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Operator section with invalid operator of 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
ftype
checkExp (OpSectionRight QualName VName
op (Info StructType
op_t) Exp
e (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
_ Info ResRetType
_ SrcLoc
loc) = do
  ftype <- SrcLoc -> QualName VName -> StructType -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
op StructType
op_t
  e' <- checkExp e
  case ftype of
    Scalar (Arrow NoMode
_ PName
m1 Diet
d1 StructType
t1 (RetType [] (Scalar (Arrow Freshness
_ PName
m2 Diet
d2 StructType
t2 (RetType [VName]
dims2 ResType
ret))))) -> do
      (t2', arrow', argext, _) <-
        SrcLoc
-> ApplyOp
-> StructType
-> Exp
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply
          SrcLoc
loc
          (QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
op, Level
1)
          (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
$ NoMode
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoMode
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow NoMode
forall a. Monoid a => a
mempty PName
m2 Diet
d2 StructType
t2 (ResRetType -> ScalarTypeBase Exp NoMode)
-> ResRetType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp Freshness
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow Freshness
Nonfresh PName
m1 Diet
d1 StructType
t1 (ResRetType -> ScalarTypeBase Exp Freshness)
-> ResRetType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims2 ResType
ret)
          Exp
e'
      case arrow' of
        Scalar (Arrow NoMode
_ PName
_ Diet
_ StructType
t1' (RetType [VName]
dims2' ResType
ret')) ->
          Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermTypeM Exp) -> Exp -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
            QualName VName
-> Info StructType
-> Exp
-> (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
QualName vn
-> f StructType
-> ExpBase f vn
-> (f (PName, ParamType), f (PName, ParamType, Maybe VName))
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
OpSectionRight
              QualName VName
op
              (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ftype)
              Exp
e'
              ((PName, ParamType) -> Info (PName, ParamType)
forall a. a -> Info a
Info (PName
m1, Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
d1 StructType
t1'), (PName, ParamType, Maybe VName)
-> Info (PName, ParamType, Maybe VName)
forall a. a -> Info a
Info (PName
m2, Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
d2 StructType
t2', Maybe VName
argext))
              (ResRetType -> Info ResRetType
forall a. a -> Info a
Info (ResRetType -> Info ResRetType) -> ResRetType -> Info ResRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims2' ResType
ret')
              SrcLoc
loc
        StructType
_ -> [Char] -> TermTypeM Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM Exp) -> [Char] -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"OpSectionRight: impossible type\n" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
arrow'
    StructType
_ ->
      SrcLoc -> Notes -> Doc () -> TermTypeM 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 SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM Exp) -> Doc () -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Operator section with invalid operator of 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
ftype
checkExp (UpdateSection [UpdateStep Info VName]
steps (Info StructType
ft) SrcLoc
loc) = do
  -- The unsized type checker has already determined the type of the
  -- parameter; we just have to instantiate its sizes. The result
  -- type is then computed by walking the steps, such that its sizes
  -- are those of the corresponding components of the parameter type.
  a <- case StructType
ft of
    Scalar (Arrow NoMode
_ PName
_ Diet
_ StructType
pt ResRetType
_) -> SrcLoc -> StructType -> TermTypeM StructType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVars SrcLoc
loc StructType
pt
    StructType
_ -> [Char] -> TermTypeM StructType
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM StructType) -> [Char] -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ [Char]
"checkExp UpdateSection: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
ft
  (steps', b, retext) <- checkSectionSteps a steps
  let ft' = 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
$ NoMode
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoMode
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow NoMode
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe StructType
a (ResRetType -> ScalarTypeBase Exp NoMode)
-> ResRetType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
retext (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
b
  pure $ UpdateSection steps' (Info ft') loc
  where
    checkSectionSteps :: TypeBase Exp o
-> [UpdateStep Info VName]
-> TermTypeM ([UpdateStep Info VName], TypeBase Exp o, [VName])
checkSectionSteps TypeBase Exp o
t [] =
      ([UpdateStep Info VName], TypeBase Exp o, [VName])
-> TermTypeM ([UpdateStep Info VName], TypeBase Exp o, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], TypeBase Exp o
t, [])
    checkSectionSteps TypeBase Exp o
t (UpdateStep Info VName
step : [UpdateStep Info VName]
rest) =
      case UpdateStep Info VName
step of
        UpdateStepField Name
f -> do
          t' <- TypeBase Exp o -> TermTypeM (TypeBase Exp o)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully TypeBase Exp o
t
          case t' of
            Scalar (Record Map Name (TypeBase Exp o)
fs)
              | Just TypeBase Exp o
f_t <- Name -> Map Name (TypeBase Exp o) -> Maybe (TypeBase Exp o)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (TypeBase Exp o)
fs -> do
                  (rest', target_t, retext) <- TypeBase Exp o
-> [UpdateStep Info VName]
-> TermTypeM ([UpdateStep Info VName], TypeBase Exp o, [VName])
checkSectionSteps TypeBase Exp o
f_t [UpdateStep Info VName]
rest
                  pure (UpdateStepField f : rest', target_t, retext)
            TypeBase Exp o
_ ->
              [Char]
-> TermTypeM ([UpdateStep Info VName], TypeBase Exp o, [VName])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> TermTypeM ([UpdateStep Info VName], TypeBase Exp o, [VName]))
-> [Char]
-> TermTypeM ([UpdateStep Info VName], TypeBase Exp o, [VName])
forall a b. (a -> b) -> a -> b
$
                [Char]
"checkExp UpdateSection: cannot project field "
                  [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
f
                  [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" from "
                  [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> TypeBase Exp o -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Exp o
t'
        UpdateStepSlice [DimIndex]
slice -> do
          slice' <- [DimIndex] -> TermTypeM [DimIndex]
checkSlice [DimIndex]
slice
          (t', retext) <- sliceShape Nothing slice' =<< normTypeFully t
          (rest', target_t, retext_rest) <- checkSectionSteps t' rest
          pure (UpdateStepSlice slice' : rest', target_t, retext <> retext_rest)
checkExp (AppExp (Loop [VName]
_ Pat ParamType
mergepat LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Exp
loopbody SrcLoc
loc) Info AppRes
_) = do
  ((sparams, mergepat', loopinit', form', loopbody'), appres) <-
    (Exp -> TermTypeM Exp)
-> UncheckedLoop -> SrcLoc -> TermTypeM (CheckedLoop, AppRes)
checkLoop Exp -> TermTypeM Exp
checkExp (Pat ParamType
mergepat, LoopInitBase Info VName
loopinit, LoopFormBase Info VName
form, Exp
loopbody) SrcLoc
loc
  pure $
    AppExp
      (Loop sparams mergepat' loopinit' form' loopbody' loc)
      (Info appres)
checkExp (Constr Name
name [Exp]
es (Info StructType
t) SrcLoc
loc) = do
  -- The sizes are absorbable: those of the payloads of the other
  -- constructors (and any not determined by the arguments) are
  -- adopted from the context, like the sizes of a hole. See Note
  -- [Size Inference].
  t' <- SrcLoc -> StructType -> TermTypeM StructType
forall o. SrcLoc -> TypeBase Exp o -> TermTypeM (TypeBase Exp o)
replaceTyVarsAbsorbable SrcLoc
loc StructType
t
  es' <- mapM checkExp es
  case t' of
    Scalar (Sum Map Name [StructType]
cs)
      | Just [StructType]
name_ts <- Name -> Map Name [StructType] -> Maybe [StructType]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
name Map Name [StructType]
cs ->
          (StructType -> StructType -> TermTypeM ())
-> [StructType] -> [StructType] -> TermTypeM ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (Usage -> StructType -> StructType -> TermTypeM ())
-> Usage -> StructType -> StructType -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"inferred variant") [StructType]
name_ts ([StructType] -> TermTypeM ()) -> [StructType] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
            (Exp -> StructType) -> [Exp] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> StructType
typeOf [Exp]
es'
    StructType
_ ->
      [Char] -> TermTypeM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM ()) -> [Char] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ [Char]
"checkExp Constr: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t'
  pure $ Constr name es' (Info t') loc
checkExp (AppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) Info AppRes
_) = do
  e1' <- Exp -> TermTypeM Exp
checkExp Exp
e1
  e2' <- checkExp e2
  e3' <- checkExp e3
  (t, retext) <- unifyBranches loc e2' e3'
  pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes t retext)
checkExp (AppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
_) = do
  e' <- Exp -> TermTypeM Exp
checkExp Exp
e
  mt <- expType e'
  (cs', t, retext) <- checkCases mt cs

  pure $ AppExp (Match e' cs' loc) (Info $ AppRes t retext)
checkExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) =
  AttrInfo VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr (AttrInfo VName -> Exp -> SrcLoc -> Exp)
-> TermTypeM (AttrInfo VName) -> TermTypeM (Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AttrInfo VName -> TermTypeM (AttrInfo VName)
forall (m :: * -> *).
MonadTypeChecker m =>
AttrInfo VName -> m (AttrInfo VName)
checkAttr AttrInfo VName
info TermTypeM (Exp -> SrcLoc -> Exp)
-> TermTypeM Exp -> TermTypeM (SrcLoc -> Exp)
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> TermTypeM Exp
checkExp Exp
e TermTypeM (SrcLoc -> Exp) -> TermTypeM SrcLoc -> TermTypeM Exp
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermTypeM SrcLoc
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

checkCase ::
  StructType ->
  CaseBase Info VName ->
  TermTypeM (CaseBase Info VName, StructType, [VName])
checkCase :: StructType
-> CaseBase Info VName
-> TermTypeM (CaseBase Info VName, StructType, [VName])
checkCase StructType
mt (CasePat PatBase Info VName StructType
p Exp
e SrcLoc
loc) =
  [SizeBinder VName]
-> PatBase Info VName StructType
-> StructType
-> (Pat ParamType
    -> TermTypeM (CaseBase Info VName, StructType, [VName]))
-> TermTypeM (CaseBase Info VName, StructType, [VName])
forall o a.
[SizeBinder VName]
-> Pat (TypeBase Exp o)
-> StructType
-> (Pat ParamType -> TermTypeM a)
-> TermTypeM a
bindingPat [] PatBase Info VName StructType
p StructType
mt ((Pat ParamType
  -> TermTypeM (CaseBase Info VName, StructType, [VName]))
 -> TermTypeM (CaseBase Info VName, StructType, [VName]))
-> (Pat ParamType
    -> TermTypeM (CaseBase Info VName, StructType, [VName]))
-> TermTypeM (CaseBase Info VName, StructType, [VName])
forall a b. (a -> b) -> a -> b
$ \Pat ParamType
p' -> do
    e' <- Exp -> TermTypeM Exp
checkExp Exp
e
    e_t <- expTypeFully e'
    (e_t', retext) <- unscopeType loc (patNames p') e_t
    pure (CasePat (fmap toStruct p') e' loc, e_t', retext)

updateFieldPath ::
  (Pretty a, Located a) =>
  a ->
  [Name] ->
  StructType ->
  StructType ->
  TermTypeM StructType
updateFieldPath :: forall a.
(Pretty a, Located a) =>
a -> [Name] -> StructType -> StructType -> TermTypeM StructType
updateFieldPath a
src [Name]
all_fs StructType
ve_t = [Name] -> [Name] -> StructType -> TermTypeM StructType
recurse [] [Name]
all_fs
  where
    recurse :: [Name] -> [Name] -> StructType -> TermTypeM StructType
recurse [Name]
seen [] StructType
t = do
      (t', _) <- Usage
-> Rigidity
-> Name
-> StructType
-> TermTypeM (StructType, Map VName Exp)
forall (m :: * -> *) d o.
MonadUnify m =>
Usage
-> Rigidity
-> Name
-> TypeBase d o
-> m (TypeBase Exp o, Map VName d)
allDimsFreshInType Usage
usage Rigidity
Nonrigid Name
"any" StructType
t
      onFailure (CheckingRecordUpdate seen t' ve_t) $
        unify usage t' ve_t
      pure ve_t
      where
        usage :: Usage
usage = Loc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage (a -> Loc
forall a. Located a => a -> Loc
locOf a
src) Text
"record update"
    recurse [Name]
seen (Name
f : [Name]
fs) (Scalar (Record Map Name StructType
m))
      | Just StructType
f_t <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name StructType
m = do
          f_t' <- [Name] -> [Name] -> StructType -> TermTypeM StructType
recurse ([Name]
seen [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ [Name
f]) [Name]
fs StructType
f_t
          pure $ Scalar $ Record $ M.insert f f_t' m
    recurse [Name]
_ [Name]
_ StructType
_ =
      Loc -> Notes -> Doc () -> TermTypeM StructType
forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError (a -> Loc
forall a. Located a => a -> Loc
locOf a
src) Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM StructType)
-> (Doc () -> Doc ()) -> Doc () -> TermTypeM StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"record-type-not-known" (Doc () -> TermTypeM StructType) -> Doc () -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Full type of"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 (a -> Doc ()
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
src)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
" is not known at this point.  Add a type annotation to the original record to disambiguate."

checkUpdateSteps ::
  SrcLoc ->
  StructType ->
  [UpdateStep Info VName] ->
  TermTypeM ([UpdateStep Info VName], StructType)
checkUpdateSteps :: SrcLoc
-> StructType
-> [UpdateStep Info VName]
-> TermTypeM ([UpdateStep Info VName], StructType)
checkUpdateSteps SrcLoc
_ StructType
t [] =
  ([UpdateStep Info VName], StructType)
-> TermTypeM ([UpdateStep Info VName], StructType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], StructType
t)
checkUpdateSteps SrcLoc
loc StructType
t (UpdateStep Info VName
step : [UpdateStep Info VName]
rest) =
  case UpdateStep Info VName
step of
    UpdateStepSlice [DimIndex]
slice -> do
      slice' <- [DimIndex] -> TermTypeM [DimIndex]
checkSlice [DimIndex]
slice
      (elem_t, _) <- sliceShape (Just (loc, Nonrigid)) slice' =<< normTypeFully t
      (rest', target_t) <- checkUpdateSteps loc elem_t rest
      pure (UpdateStepSlice slice' : rest', target_t)
    UpdateStepField Name
f -> do
      t' <- StructType -> TermTypeM StructType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully StructType
t
      case t' of
        Scalar (Record Map Name StructType
fs) | Just StructType
f_t <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name StructType
fs -> do
          (rest', target_t) <- SrcLoc
-> StructType
-> [UpdateStep Info VName]
-> TermTypeM ([UpdateStep Info VName], StructType)
checkUpdateSteps SrcLoc
loc StructType
f_t [UpdateStep Info VName]
rest
          pure (UpdateStepField f : rest', target_t)
        StructType
_ -> [Char] -> TermTypeM ([UpdateStep Info VName], StructType)
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM ([UpdateStep Info VName], StructType))
-> [Char] -> TermTypeM ([UpdateStep Info VName], StructType)
forall a b. (a -> b) -> a -> b
$ [Char]
"checkUpdateSteps: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Show a => a -> [Char]
show StructType
t'

checkCases ::
  StructType ->
  NE.NonEmpty (CaseBase Info VName) ->
  TermTypeM (NE.NonEmpty (CaseBase Info VName), StructType, [VName])
checkCases :: StructType
-> NonEmpty (CaseBase Info VName)
-> TermTypeM (NonEmpty (CaseBase Info VName), StructType, [VName])
checkCases StructType
mt NonEmpty (CaseBase Info VName)
rest_cs =
  case NonEmpty (CaseBase Info VName)
-> (CaseBase Info VName, Maybe (NonEmpty (CaseBase Info VName)))
forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (CaseBase Info VName)
rest_cs of
    (CaseBase Info VName
c, Maybe (NonEmpty (CaseBase Info VName))
Nothing) -> do
      (c', t, retext) <- StructType
-> CaseBase Info VName
-> TermTypeM (CaseBase Info VName, StructType, [VName])
checkCase StructType
mt CaseBase Info VName
c
      pure (NE.singleton c', t, retext)
    (CaseBase Info VName
c, Just NonEmpty (CaseBase Info VName)
cs) -> do
      ((c', c_t, _), (cs', cs_t, _)) <-
        (,) ((CaseBase Info VName, StructType, [VName])
 -> (NonEmpty (CaseBase Info VName), StructType, [VName])
 -> ((CaseBase Info VName, StructType, [VName]),
     (NonEmpty (CaseBase Info VName), StructType, [VName])))
-> TermTypeM (CaseBase Info VName, StructType, [VName])
-> TermTypeM
     ((NonEmpty (CaseBase Info VName), StructType, [VName])
      -> ((CaseBase Info VName, StructType, [VName]),
          (NonEmpty (CaseBase Info VName), StructType, [VName])))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType
-> CaseBase Info VName
-> TermTypeM (CaseBase Info VName, StructType, [VName])
checkCase StructType
mt CaseBase Info VName
c TermTypeM
  ((NonEmpty (CaseBase Info VName), StructType, [VName])
   -> ((CaseBase Info VName, StructType, [VName]),
       (NonEmpty (CaseBase Info VName), StructType, [VName])))
-> TermTypeM (NonEmpty (CaseBase Info VName), StructType, [VName])
-> TermTypeM
     ((CaseBase Info VName, StructType, [VName]),
      (NonEmpty (CaseBase Info VName), StructType, [VName]))
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> StructType
-> NonEmpty (CaseBase Info VName)
-> TermTypeM (NonEmpty (CaseBase Info VName), StructType, [VName])
checkCases StructType
mt NonEmpty (CaseBase Info VName)
cs
      (brancht, retext) <- unifyBranchTypes (srclocOf c) c_t cs_t
      pure (NE.cons c' cs', brancht, retext)

checkSlice :: SliceBase Info VName -> TermTypeM [DimIndex]
checkSlice :: [DimIndex] -> TermTypeM [DimIndex]
checkSlice = (DimIndex -> TermTypeM DimIndex)
-> [DimIndex] -> TermTypeM [DimIndex]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM DimIndex -> TermTypeM DimIndex
checkDimIndex
  where
    checkDimIndex :: DimIndex -> TermTypeM DimIndex
checkDimIndex (DimFix Exp
i) =
      Exp -> DimIndex
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (Exp -> DimIndex) -> TermTypeM Exp -> TermTypeM DimIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> TermTypeM Exp
checkExp Exp
i
    checkDimIndex (DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
s) =
      Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndex
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndex)
-> TermTypeM (Maybe Exp)
-> TermTypeM (Maybe Exp -> Maybe Exp -> DimIndex)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> TermTypeM Exp) -> Maybe Exp -> TermTypeM (Maybe Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Exp -> TermTypeM Exp
checkExp Maybe Exp
i TermTypeM (Maybe Exp -> Maybe Exp -> DimIndex)
-> TermTypeM (Maybe Exp) -> TermTypeM (Maybe Exp -> DimIndex)
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Exp -> TermTypeM Exp) -> Maybe Exp -> TermTypeM (Maybe Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Exp -> TermTypeM Exp
checkExp Maybe Exp
j TermTypeM (Maybe Exp -> DimIndex)
-> TermTypeM (Maybe Exp) -> TermTypeM DimIndex
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Exp -> TermTypeM Exp) -> Maybe Exp -> TermTypeM (Maybe Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Exp -> TermTypeM Exp
checkExp Maybe Exp
s

instantiateDimsInReturnType ::
  SrcLoc ->
  Maybe (QualName VName) ->
  ResRetType ->
  TermTypeM (ResType, [VName])
instantiateDimsInReturnType :: SrcLoc
-> Maybe (QualName VName)
-> ResRetType
-> TermTypeM (ResType, [VName])
instantiateDimsInReturnType SrcLoc
loc Maybe (QualName VName)
fname (RetType [VName]
dims ResType
t)
  | [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
dims =
      (ResType, [VName]) -> TermTypeM (ResType, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResType
t, [VName]
forall a. Monoid a => a
mempty)
  | Bool
otherwise = do
      dims' <- (VName -> TermTypeM VName) -> [VName] -> TermTypeM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> TermTypeM VName
new [VName]
dims
      pure (first (onDim $ zip dims $ map (ExpSubst . (`sizeFromName` loc) . qualName) dims') t, dims')
  where
    new :: VName -> TermTypeM VName
new =
      SrcLoc -> RigidSource -> Name -> TermTypeM VName
forall a. Located a => a -> RigidSource -> Name -> TermTypeM VName
forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim SrcLoc
loc (Maybe (QualName VName) -> RigidSource
RigidRet Maybe (QualName VName)
fname)
        (Name -> TermTypeM VName)
-> (VName -> Name) -> VName -> TermTypeM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Name
nameFromText
        (Text -> Name) -> (VName -> Text) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Text -> Text
T.takeWhile Char -> Bool
isAscii
        (Text -> Text) -> (VName -> Text) -> VName -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Text
baseText
    onDim :: [(VName, Subst StructRetType)] -> a -> a
onDim [(VName, Subst StructRetType)]
dims' = TypeSubs -> a -> a
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> [(VName, Subst StructRetType)] -> Maybe (Subst StructRetType)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(VName, Subst StructRetType)]
dims')

-- Some information about the function/operator we are trying to
-- apply, and how many arguments it has previously accepted.  Used for
-- generating nicer type errors.
type ApplyOp = (Maybe (QualName VName), Int)

-- | Extract all those names that are bound inside the type.
boundInsideType :: TypeBase Size o -> S.Set VName
boundInsideType :: forall o. TypeBase Exp o -> Set VName
boundInsideType (Array o
_ Shape Exp
_ ScalarTypeBase Exp NoMode
t) = StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase Exp NoMode
t)
boundInsideType (Scalar Prim {}) = Set VName
forall a. Monoid a => a
mempty
boundInsideType (Scalar (TypeVar o
_ QualName VName
_ [TypeArg Exp]
targs)) = (TypeArg Exp -> Set VName) -> [TypeArg Exp] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeArg Exp -> Set VName
f [TypeArg Exp]
targs
  where
    f :: TypeArg Exp -> Set VName
f (TypeArgType StructType
t) = StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType StructType
t
    f TypeArgDim {} = Set VName
forall a. Monoid a => a
mempty
boundInsideType (Scalar (Record Map Name (TypeBase Exp o)
fs)) = (TypeBase Exp o -> Set VName)
-> Map Name (TypeBase Exp o) -> Set VName
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase Exp o -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType Map Name (TypeBase Exp o)
fs
boundInsideType (Scalar (Sum Map Name [TypeBase Exp o]
cs)) = ([TypeBase Exp o] -> Set VName)
-> Map Name [TypeBase Exp o] -> Set VName
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((TypeBase Exp o -> Set VName) -> [TypeBase Exp o] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase Exp o -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType) Map Name [TypeBase Exp o]
cs
boundInsideType (Scalar (Arrow o
_ PName
pn Diet
_ StructType
t1 (RetType [VName]
dims ResType
t2))) =
  Set VName
pn' Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType StructType
t1 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]
dims Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> ResType -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType ResType
t2
  where
    pn' :: Set VName
pn' = case PName
pn of
      PName
Unnamed -> Set VName
forall a. Monoid a => a
mempty
      Named VName
v -> VName -> Set VName
forall a. a -> Set a
S.singleton VName
v

-- Returns the sizes of the immediate type produced,
-- the sizes of parameter types, and the sizes of return types.
dimUses :: TypeBase Size o -> (Names, Names)
dimUses :: forall o. TypeBase Exp o -> (Set VName, Set VName)
dimUses = (State (Set VName, Set VName) (TypeBase () o)
 -> (Set VName, Set VName) -> (Set VName, Set VName))
-> (Set VName, Set VName)
-> State (Set VName, Set VName) (TypeBase () o)
-> (Set VName, Set VName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (Set VName, Set VName) (TypeBase () o)
-> (Set VName, Set VName) -> (Set VName, Set VName)
forall s a. State s a -> s -> s
execState (Set VName, Set VName)
forall a. Monoid a => a
mempty (State (Set VName, Set VName) (TypeBase () o)
 -> (Set VName, Set VName))
-> (TypeBase Exp o -> State (Set VName, Set VName) (TypeBase () o))
-> TypeBase Exp o
-> (Set VName, Set VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName
 -> DimPos -> Exp -> StateT (Set VName, Set VName) Identity ())
-> TypeBase Exp o -> State (Set VName, Set VName) (TypeBase () o)
forall (f :: * -> *) fdim tdim o.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim o -> f (TypeBase tdim o)
traverseDims Set VName
-> DimPos -> Exp -> StateT (Set VName, Set VName) Identity ()
forall {m :: * -> *}.
MonadState (Set VName, Set VName) m =>
Set VName -> DimPos -> Exp -> m ()
f
  where
    f :: Set VName -> DimPos -> Exp -> m ()
f Set VName
bound DimPos
pos Exp
e =
      case DimPos
pos of
        DimPos
PosImmediate ->
          ((Set VName, Set VName) -> (Set VName, Set VName)) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> Set VName
fvVars FV
fv, Set VName
forall a. Monoid a => a
mempty) (Set VName, Set VName)
-> (Set VName, Set VName) -> (Set VName, Set VName)
forall a. Semigroup a => a -> a -> a
<>)
        DimPos
PosParam ->
          ((Set VName, Set VName) -> (Set VName, Set VName)) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName
forall a. Monoid a => a
mempty, FV -> Set VName
fvVars FV
fv) (Set VName, Set VName)
-> (Set VName, Set VName) -> (Set VName, Set VName)
forall a. Semigroup a => a -> a -> a
<>)
        DimPos
PosReturn -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      where
        fv :: FV
fv = Exp -> FV
freeInExp Exp
e FV -> Set VName -> FV
`freeWithout` Set VName
bound

checkApply ::
  SrcLoc ->
  ApplyOp ->
  StructType ->
  Exp ->
  TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply :: SrcLoc
-> ApplyOp
-> StructType
-> Exp
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (Maybe (QualName VName)
fname, Level
_) (Scalar (Arrow NoMode
_ PName
pname Diet
_ StructType
tp1 ResRetType
tp2)) Exp
argexp = do
  let argtype :: StructType
argtype = Exp -> StructType
typeOf Exp
argexp
  Checking
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (Maybe (QualName VName)
-> Exp -> StructType -> StructType -> Checking
CheckingApply Maybe (QualName VName)
fname Exp
argexp StructType
tp1 StructType
argtype) (TermTypeM (StructType, StructType, Maybe VName, [VName])
 -> TermTypeM (StructType, StructType, Maybe VName, [VName]))
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
forall a b. (a -> b) -> a -> b
$ do
    Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (Exp -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Exp
argexp Text
"use as function argument") StructType
tp1 StructType
argtype

    -- Perform substitutions of instantiated variables in the types.
    (tp2_inst, ext) <- SrcLoc
-> Maybe (QualName VName)
-> ResRetType
-> TermTypeM (ResType, [VName])
instantiateDimsInReturnType SrcLoc
loc Maybe (QualName VName)
fname (ResRetType -> TermTypeM (ResType, [VName]))
-> TermTypeM ResRetType -> TermTypeM (ResType, [VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ResRetType -> TermTypeM ResRetType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully ResRetType
tp2
    argtype' <- normTypeFully argtype

    -- Unification against the argument type may have determined that
    -- some instantiated sizes are existential. Their occurrences in
    -- the return type are replaced with fresh rigid sizes, bound at
    -- the innermost possible position; those bound at the top level
    -- become existentials of the application. The pending size
    -- variables themselves are left alone; occurrences of them
    -- remaining in the AST are existentially bound by
    -- 'bindExistentialInsts' at the end.
    constraints <- getConstraints
    let (inst_pending, inst_reps) = pendingInstSizes constraints
        repOf VName
v = Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> (VName -> Exp) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName (SrcLoc -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf SrcLoc
loc) (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> Subst t) -> Maybe VName -> Maybe (Subst t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
inst_reps
        tp2_subst = TypeSubs -> ResType -> ResType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
forall {t}. VName -> Maybe (Subst t)
repOf ResType
tp2_inst
    (tp2', inst_ext) <-
      -- 'sizeFree' can only change the type if a pending instantiated size
      -- occurs free in it, so check for a fast path.
      if any inst_pending (fvVars (freeInType tp2_subst))
        then sizeFree loc (find inst_pending . fvVars . freeInExp) tp2_subst
        else pure (tp2_subst, [])
    let ext' = [VName]
ext [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
inst_ext

    -- Check whether this would produce an impossible return type.
    let (tp2_produced_dims, tp2_paramdims) = dimUses tp2'
        problematic = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext' Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
boundInsideType StructType
argtype'
        problem = (VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
problematic) (Set VName
tp2_paramdims Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Set VName
tp2_produced_dims)
    when (not (S.null problematic) && problem) $ do
      typeError loc mempty . withIndexLink "existential-param-ret" $
        "Existential size would appear in function parameter of return type:"
          </> indent 2 (pretty (RetType ext' tp2'))
          </> textwrap "This is usually because a higher-order function is used with functional arguments that return existential sizes or locally named sizes, which are then used as parameters of other function arguments."

    (argext, tp2'') <-
      case pname of
        Named VName
pname'
          | VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member VName
pname' (FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ ResType -> FV
forall o. TypeBase Exp o -> FV
freeInType ResType
tp2') ->
              if Exp -> Bool
hasBinding Exp
argexp
                then do
                  SrcLoc -> Doc () -> TermTypeM ()
forall loc. Located loc => loc -> Doc () -> TermTypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
argexp) (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
                    Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink
                      Doc ()
"size-expression-bind"
                      Doc ()
"Size expression with binding is replaced by unknown size."
                  d <- Exp -> RigidSource -> Name -> TermTypeM VName
forall a. Located a => a -> RigidSource -> Name -> TermTypeM VName
forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim Exp
argexp (Maybe (QualName VName) -> Text -> RigidSource
RigidArg Maybe (QualName VName)
fname (Text -> RigidSource) -> Text -> RigidSource
forall a b. (a -> b) -> a -> b
$ ExpBase NoInfo VName -> Text
forall a. Pretty a => a -> Text
prettyTextOneLine (ExpBase NoInfo VName -> Text) -> ExpBase NoInfo VName -> Text
forall a b. (a -> b) -> a -> b
$ Exp -> ExpBase NoInfo VName
bareExp Exp
argexp) Name
"n"
                  let parsubst VName
v =
                        if VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
pname'
                          then Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> Exp -> Subst t
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
d) (SrcLoc -> Exp) -> SrcLoc -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
argexp
                          else Maybe (Subst t)
forall a. Maybe a
Nothing
                  pure (Just d, applySubst parsubst $ toStruct tp2')
                else
                  let parsubst :: VName -> Maybe (Subst t)
parsubst VName
v =
                        if VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
pname'
                          then Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> Exp -> Subst t
forall a b. (a -> b) -> a -> b
$ Exp -> Maybe Exp -> Exp
forall a. a -> Maybe a -> a
fromMaybe Exp
argexp (Maybe Exp -> Exp) -> Maybe Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Maybe Exp
stripExp Exp
argexp
                          else Maybe (Subst t)
forall a. Maybe a
Nothing
                   in (Maybe VName, StructType) -> TermTypeM (Maybe VName, StructType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe VName
forall a. Maybe a
Nothing, TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
forall {t}. VName -> Maybe (Subst t)
parsubst (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
tp2')
        PName
_ -> (Maybe VName, StructType) -> TermTypeM (Maybe VName, StructType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe VName
forall a. Maybe a
Nothing, ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
tp2')

    pure (tp1, tp2'', argext, ext')
checkApply SrcLoc
_ ApplyOp
_ StructType
_ Exp
_ =
  [Char] -> TermTypeM (StructType, StructType, Maybe VName, [VName])
forall a. HasCallStack => [Char] -> a
error [Char]
"checkApply: array"

-- | Type-check a single expression in isolation.  This expression may
-- turn out to be polymorphic, in which case the list of type
-- parameters will be non-empty.
checkOneExp :: ExpBase NoInfo VName -> TypeM ([TypeParam], Exp)
checkOneExp :: ExpBase NoInfo VName -> TypeM ([TypeParamBase VName], Exp)
checkOneExp ExpBase NoInfo VName
e = do
  (maybe_tysubsts, e') <- ExpBase NoInfo VName
-> TypeM
     (Either
        TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode)),
      Exp)
Unsized.checkSingleExp ExpBase NoInfo VName
e
  case maybe_tysubsts of
    Left TypeError
err -> TypeError -> TypeM ([TypeParamBase VName], Exp)
forall a. TypeError -> TypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError TypeError
err
    Right ([TypeParamBase VName]
generalised, Map VName (TypeBase () NoMode)
tysubsts) -> (Exp -> TermTypeM Exp)
-> Map VName (TypeBase () NoMode)
-> TermTypeM ([TypeParamBase VName], Exp)
-> TypeM ([TypeParamBase VName], Exp)
forall a.
(Exp -> TermTypeM Exp)
-> Map VName (TypeBase () NoMode) -> TermTypeM a -> TypeM a
runTermTypeM Exp -> TermTypeM Exp
checkExp Map VName (TypeBase () NoMode)
tysubsts (TermTypeM ([TypeParamBase VName], Exp)
 -> TypeM ([TypeParamBase VName], Exp))
-> TermTypeM ([TypeParamBase VName], Exp)
-> TypeM ([TypeParamBase VName], Exp)
forall a b. (a -> b) -> a -> b
$ do
      e'' <- Exp -> TermTypeM Exp
checkExp Exp
e'
      let t = Exp -> StructType
typeOf Exp
e''
      (tparams, _, _) <-
        letGeneralise (nameFromString "<exp>") (srclocOf e) generalised [] $ toRes Nonfresh t
      detectAmbiguousSizes
      e''' <- bindExistentialInsts =<< normTypeFully e''
      localChecks tparams e'''
      causalityCheck e'''
      pure (tparams, e''')

-- | Type-check a single size expression in isolation.  This expression may
-- turn out to be polymorphic, in which case it is unified with i64.
checkSizeExp :: ExpBase NoInfo VName -> TypeM Exp
checkSizeExp :: ExpBase NoInfo VName -> TypeM Exp
checkSizeExp ExpBase NoInfo VName
e = do
  (maybe_tysubsts, e') <- ExpBase NoInfo VName
-> TypeM
     (Either TypeError ([UnconTyVar], Map VName (TypeBase () NoMode)),
      Exp)
Unsized.checkSizeExp ExpBase NoInfo VName
e
  case maybe_tysubsts of
    Left TypeError
err -> TypeError -> TypeM Exp
forall a. TypeError -> TypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError TypeError
err
    Right ([UnconTyVar]
_generalised, Map VName (TypeBase () NoMode)
tysubsts) -> (Exp -> TermTypeM Exp)
-> Map VName (TypeBase () NoMode) -> TermTypeM Exp -> TypeM Exp
forall a.
(Exp -> TermTypeM Exp)
-> Map VName (TypeBase () NoMode) -> TermTypeM a -> TypeM a
runTermTypeM Exp -> TermTypeM Exp
checkExp Map VName (TypeBase () NoMode)
tysubsts (TermTypeM Exp -> TypeM Exp) -> TermTypeM Exp -> TypeM Exp
forall a b. (a -> b) -> a -> b
$ do
      e'' <- Exp -> TermTypeM Exp
checkExp Exp
e'
      when (hasBinding e'') $
        typeError (srclocOf e'') mempty . withIndexLink "size-expression-bind" $
          "Size expression with binding is forbidden."
      normTypeFully e''

-- Verify that all sum type constructors and empty array literals have
-- a size that is known (rigid or a type parameter).  This is to
-- ensure that we can actually determine their shape at run-time.
causalityCheck :: Exp -> TermTypeM ()
causalityCheck :: Exp -> TermTypeM ()
causalityCheck Exp
binding_body = do
  constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints

  let checkCausality Doc ()
what Set VName
known TypeBase Exp o
t a
loc
        | (VName
d, Loc
dloc) : [(VName, Loc)]
_ <-
            (VName -> Maybe (VName, Loc)) -> [VName] -> [(VName, Loc)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Constraints -> Set VName -> VName -> Maybe (VName, Loc)
forall {a} {a}.
Ord a =>
Map a (a, Constraint) -> Set a -> a -> Maybe (a, Loc)
unknown Constraints
constraints Set VName
known) ([VName] -> [(VName, Loc)]) -> [VName] -> [(VName, Loc)]
forall a b. (a -> b) -> a -> b
$
              Set VName -> [VName]
forall a. Set a -> [a]
S.toList (FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> FV
forall o. TypeBase Exp o -> FV
freeInType TypeBase Exp o
t) =
            t (Either TypeError) a -> Maybe (t (Either TypeError) a)
forall a. a -> Maybe a
Just (t (Either TypeError) a -> Maybe (t (Either TypeError) a))
-> t (Either TypeError) a -> Maybe (t (Either TypeError) a)
forall a b. (a -> b) -> a -> b
$ Either TypeError a -> t (Either TypeError) a
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either TypeError a -> t (Either TypeError) a)
-> Either TypeError a -> t (Either TypeError) a
forall a b. (a -> b) -> a -> b
$ Doc ()
-> Loc -> VName -> Loc -> TypeBase Exp o -> Either TypeError a
forall {v} {a} {b} {b}.
(IsName v, Pretty a, Located b) =>
Doc () -> Loc -> v -> b -> a -> Either TypeError b
causality Doc ()
what (a -> Loc
forall a. Located a => a -> Loc
locOf a
loc) VName
d Loc
dloc TypeBase Exp o
t
        | Bool
otherwise = Maybe (t (Either TypeError) a)
forall a. Maybe a
Nothing

      checkParamCausality Set VName
known Pat (TypeBase Exp o)
p =
        Doc ()
-> Set VName
-> TypeBase Exp o
-> Loc
-> Maybe (t (Either TypeError) a)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality (Pat (TypeBase Exp o) -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Pat (TypeBase Exp o) -> Doc ann
pretty Pat (TypeBase Exp o)
p) Set VName
known (Pat (TypeBase Exp o) -> TypeBase Exp o
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType Pat (TypeBase Exp o)
p) (Pat (TypeBase Exp o) -> Loc
forall a. Located a => a -> Loc
locOf Pat (TypeBase Exp o)
p)

      collectingNewKnown = Either TypeError (Set VName)
-> StateT (Set VName) (Either TypeError) (Set VName)
forall (m :: * -> *) a. Monad m => m a -> StateT (Set VName) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either TypeError (Set VName)
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> (StateT (Set VName) (Either TypeError) a
    -> Either TypeError (Set VName))
-> StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StateT (Set VName) (Either TypeError) a
 -> Set VName -> Either TypeError (Set VName))
-> Set VName
-> StateT (Set VName) (Either TypeError) a
-> Either TypeError (Set VName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (Set VName) (Either TypeError) a
-> Set VName -> Either TypeError (Set VName)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT Set VName
forall a. Monoid a => a
mempty

      onExp ::
        S.Set VName ->
        Exp ->
        StateT (S.Set VName) (Either TypeError) Exp

      onExp Set VName
known (Var QualName VName
v (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Exp
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality (Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (QualName VName -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName VName -> Doc ann
pretty QualName VName
v)) Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Exp
bad
      onExp Set VName
known (UpdateSection [UpdateStep Info VName]
_ (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Exp
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"projection section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Exp
bad
      onExp Set VName
known (OpSectionRight QualName VName
_ (Info StructType
t) Exp
_ (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
_ Info ResRetType
_ SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Exp
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"operator section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Exp
bad
      onExp Set VName
known (OpSectionLeft QualName VName
_ (Info StructType
t) Exp
_ (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
_ (Info ResRetType, Info [VName])
_ SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Exp
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"operator section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Exp
bad
      onExp Set VName
known (ArrayLit [] (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Exp
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"empty array" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Exp
bad
      onExp Set VName
known (Hole (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Exp
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"hole" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Exp
bad
      onExp Set VName
known e :: Exp
e@(Lambda [Pat ParamType]
params Exp
body Maybe (TypeExp Exp VName)
_ Info ResRetType
_ SrcLoc
_)
        | StateT (Set VName) (Either TypeError) Exp
bad : [StateT (Set VName) (Either TypeError) Exp]
_ <- (Pat ParamType
 -> Maybe (StateT (Set VName) (Either TypeError) Exp))
-> [Pat ParamType] -> [StateT (Set VName) (Either TypeError) Exp]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Set VName
-> Pat ParamType
-> Maybe (StateT (Set VName) (Either TypeError) Exp)
forall {t :: (* -> *) -> * -> *} {o} {a}.
(MonadTrans t, Pretty o) =>
Set VName -> Pat (TypeBase Exp o) -> Maybe (t (Either TypeError) a)
checkParamCausality Set VName
known) [Pat ParamType]
params =
            StateT (Set VName) (Either TypeError) Exp
bad
        | Bool
otherwise = do
            -- Existentials coming into existence in the lambda body
            -- are not known outside of it.
            StateT (Set VName) (Either TypeError) (Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) (Set VName)
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) (Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall {a}.
StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
collectingNewKnown (StateT (Set VName) (Either TypeError) Exp
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Exp -> StateT (Set VName) (Either TypeError) Exp
onExp Set VName
known Exp
body
            Exp -> StateT (Set VName) (Either TypeError) Exp
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
      onExp Set VName
known e :: Exp
e@(AppExp (LetPat [SizeBinder VName]
_ PatBase Info VName StructType
_ Exp
bindee_e Exp
body_e SrcLoc
_) (Info AppRes
res)) = do
        Set VName
-> Exp
-> Exp
-> [VName]
-> StateT (Set VName) (Either TypeError) ()
sequencePoint Set VName
known Exp
bindee_e Exp
body_e ([VName] -> StateT (Set VName) (Either TypeError) ())
-> [VName] -> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ AppRes -> [VName]
appResExt AppRes
res
        Exp -> StateT (Set VName) (Either TypeError) Exp
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
      onExp Set VName
known e :: Exp
e@(AppExp (Match Exp
scrutinee NonEmpty (CaseBase Info VName)
cs SrcLoc
_) (Info AppRes
res)) = do
        new_known <- StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall {a}.
StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
collectingNewKnown (StateT (Set VName) (Either TypeError) Exp
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Exp -> StateT (Set VName) (Either TypeError) Exp
onExp Set VName
known Exp
scrutinee
        void $ recurse (new_known <> known) cs
        modify ((new_known <> S.fromList (appResExt res)) <>)
        pure e
      onExp Set VName
known e :: Exp
e@(AppExp (Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
_) (Info AppRes
res)) = do
        Set VName
-> [(Info (Maybe VName), Exp)]
-> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known ([(Info (Maybe VName), Exp)]
 -> StateT (Set VName) (Either TypeError) ())
-> [(Info (Maybe VName), Exp)]
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ [(Info (Maybe VName), Exp)] -> [(Info (Maybe VName), Exp)]
forall a. [a] -> [a]
reverse ([(Info (Maybe VName), Exp)] -> [(Info (Maybe VName), Exp)])
-> [(Info (Maybe VName), Exp)] -> [(Info (Maybe VName), Exp)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args
        Exp -> StateT (Set VName) (Either TypeError) Exp
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
        where
          seqArgs :: Set VName
-> [(Info (Maybe VName), Exp)]
-> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known' [] = do
            StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) Exp
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName -> Exp -> StateT (Set VName) (Either TypeError) Exp
onExp Set VName
known' Exp
f
            (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (AppRes -> [VName]
appResExt AppRes
res) Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<>)
          seqArgs Set VName
known' ((Info Maybe VName
p, Exp
x) : [(Info (Maybe VName), Exp)]
xs) = do
            new_known <- StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall {a}.
StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
collectingNewKnown (StateT (Set VName) (Either TypeError) Exp
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Exp -> StateT (Set VName) (Either TypeError) Exp
onExp Set VName
known' Exp
x
            void $ seqArgs (new_known <> known') xs
            modify ((new_known <> S.fromList (maybeToList p)) <>)
      onExp Set VName
known e :: Exp
e@(Constr Name
v [Exp]
args (Info StructType
t) SrcLoc
loc) = do
        Set VName -> [Exp] -> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known [Exp]
args
        Exp -> StateT (Set VName) (Either TypeError) Exp
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
        where
          seqArgs :: Set VName -> [Exp] -> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known' []
            | Just StateT (Set VName) (Either TypeError) ()
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) ())
forall {t :: (* -> *) -> * -> *} {o} {a} {a}.
(MonadTrans t, Pretty o, Located a) =>
Doc ()
-> Set VName
-> TypeBase Exp o
-> a
-> Maybe (t (Either TypeError) a)
checkCausality (Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Doc ()
"#" Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Name -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Name -> Doc ann
pretty Name
v)) Set VName
known' StructType
t SrcLoc
loc =
                StateT (Set VName) (Either TypeError) ()
bad
            | Bool
otherwise =
                () -> StateT (Set VName) (Either TypeError) ()
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          seqArgs Set VName
known' (Exp
x : [Exp]
xs) = do
            new_known <- StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall {a}.
StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
collectingNewKnown (StateT (Set VName) (Either TypeError) Exp
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Exp -> StateT (Set VName) (Either TypeError) Exp
onExp Set VName
known' Exp
x
            void $ seqArgs (new_known <> known') xs
            modify (new_known <>)
      onExp
        Set VName
known
        e :: Exp
e@(AppExp (BinOp (QualName VName
f, SrcLoc
floc) Info StructType
ft (Exp
x, Info Maybe VName
xp) (Exp
y, Info Maybe VName
yp) SrcLoc
_) (Info AppRes
res)) = do
          args_known <-
            StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) (Set VName)
forall {a}.
StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
collectingNewKnown (StateT (Set VName) (Either TypeError) ()
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName
-> Exp
-> Exp
-> [VName]
-> StateT (Set VName) (Either TypeError) ()
sequencePoint Set VName
known Exp
x Exp
y ([VName] -> StateT (Set VName) (Either TypeError) ())
-> [VName] -> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ [Maybe VName] -> [VName]
forall a. [Maybe a] -> [a]
catMaybes [Maybe VName
xp, Maybe VName
yp]
          void $ onExp (args_known <> known) (Var f ft floc)
          modify ((args_known <> S.fromList (appResExt res)) <>)
          pure e
      onExp Set VName
known e :: Exp
e@(AppExp AppExpBase Info VName
e' (Info AppRes
res)) = do
        Set VName
-> AppExpBase Info VName
-> StateT (Set VName) (Either TypeError) ()
forall {a}.
ASTMappable a =>
Set VName -> a -> StateT (Set VName) (Either TypeError) ()
recurse Set VName
known AppExpBase Info VName
e'
        (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (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 (AppRes -> [VName]
appResExt AppRes
res))
        Exp -> StateT (Set VName) (Either TypeError) Exp
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
      onExp Set VName
known Exp
e = do
        Set VName -> Exp -> StateT (Set VName) (Either TypeError) ()
forall {a}.
ASTMappable a =>
Set VName -> a -> StateT (Set VName) (Either TypeError) ()
recurse Set VName
known Exp
e
        Exp -> StateT (Set VName) (Either TypeError) Exp
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

      recurse Set VName
known = StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) a
 -> StateT (Set VName) (Either TypeError) ())
-> (a -> StateT (Set VName) (Either TypeError) a)
-> a
-> StateT (Set VName) (Either TypeError) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASTMapper (StateT (Set VName) (Either TypeError))
-> a -> StateT (Set VName) (Either TypeError) a
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> a -> m a
astMap ASTMapper (StateT (Set VName) (Either TypeError))
mapper
        where
          mapper :: ASTMapper (StateT (Set VName) (Either TypeError))
mapper = ASTMapper (StateT (Set VName) (Either TypeError))
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = onExp known}

      sequencePoint Set VName
known Exp
x Exp
y [VName]
ext = do
        new_known <- StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall {a}.
StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
collectingNewKnown (StateT (Set VName) (Either TypeError) Exp
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Exp
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Exp -> StateT (Set VName) (Either TypeError) Exp
onExp Set VName
known Exp
x
        void $ onExp (new_known <> known) y
        modify ((new_known <> S.fromList ext) <>)

  either throwError (const $ pure ()) $
    evalStateT (onExp mempty binding_body) mempty
  where
    unknown :: Map a (a, Constraint) -> Set a -> a -> Maybe (a, Loc)
unknown Map a (a, Constraint)
constraints Set a
known a
v = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ a
v a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set a
known
      loc <- case (a, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((a, Constraint) -> Constraint)
-> Maybe (a, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Map a (a, Constraint) -> Maybe (a, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
v Map a (a, Constraint)
constraints of
        Just (UnknownSize Loc
loc RigidSource
_) -> Loc -> Maybe Loc
forall a. a -> Maybe a
Just Loc
loc
        Maybe Constraint
_ -> Maybe Loc
forall a. Maybe a
Nothing
      pure (v, loc)

    causality :: Doc () -> Loc -> v -> b -> a -> Either TypeError b
causality Doc ()
what Loc
loc v
d b
dloc a
t =
      TypeError -> Either TypeError b
forall a b. a -> Either a b
Left (TypeError -> Either TypeError b)
-> (Doc () -> TypeError) -> Doc () -> Either TypeError b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Loc -> Notes -> Doc () -> TypeError
TypeError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TypeError) -> (Doc () -> Doc ()) -> Doc () -> TypeError
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"causality-check" (Doc () -> Either TypeError b) -> Doc () -> Either TypeError b
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Causality check: size"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (v -> Doc ()
forall a. v -> Doc a
forall v a. IsName v => v -> Doc a
prettyName v
d)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"needed for type of"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
what
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
forall ann. Doc ann
colon
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 (a -> Doc ()
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
t)
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"But"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (v -> Doc ()
forall a. v -> Doc a
forall v a. IsName v => v -> Doc a
prettyName v
d)
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"is computed at"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> [Char] -> Doc ()
forall ann. [Char] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Loc -> b -> [Char]
forall a b. (Located a, Located b) => a -> b -> [Char]
locStrRel Loc
loc b
dloc)
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
""
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Hint:"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
align
              ( Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
"Bind the expression producing"
                  Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (v -> Doc ()
forall a. v -> Doc a
forall v a. IsName v => v -> Doc a
prettyName v
d)
                  Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"with 'let' beforehand."
              )

mustBeIrrefutable :: (MonadTypeChecker f) => Pat StructType -> f ()
mustBeIrrefutable :: forall (f :: * -> *).
MonadTypeChecker f =>
PatBase Info VName StructType -> f ()
mustBeIrrefutable PatBase Info VName StructType
p = do
  case [PatBase Info VName StructType] -> [Match ()]
unmatched [PatBase Info VName StructType
p] of
    [] -> () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    [Match ()]
ps' ->
      PatBase Info VName StructType -> Notes -> Doc () -> f ()
forall loc a. Located loc => loc -> Notes -> Doc () -> f a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError PatBase Info VName StructType
p Notes
forall a. Monoid a => a
mempty (Doc () -> f ()) -> (Doc () -> Doc ()) -> Doc () -> f ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"refutable-pattern" (Doc () -> f ()) -> Doc () -> f ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Refutable pattern not allowed here.\nUnmatched cases:"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 ([Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
stack ((Match () -> Doc ()) -> [Match ()] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map Match () -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Match () -> Doc ann
pretty [Match ()]
ps'))

supportsEquality :: TypeBase dim o -> Bool
supportsEquality :: forall dim o. TypeBase dim o -> Bool
supportsEquality (Array o
_ Shape dim
_ ScalarTypeBase dim NoMode
t) = TypeBase dim NoMode -> Bool
forall dim o. TypeBase dim o -> Bool
supportsEquality (TypeBase dim NoMode -> Bool) -> TypeBase dim NoMode -> Bool
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim NoMode -> TypeBase dim NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase dim NoMode
t
supportsEquality (Scalar Prim {}) = Bool
True
supportsEquality (Scalar TypeVar {}) = Bool
False
supportsEquality (Scalar (Record Map Name (TypeBase dim o)
fs)) = (TypeBase dim o -> Bool) -> Map Name (TypeBase dim o) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim o -> Bool
forall dim o. TypeBase dim o -> Bool
supportsEquality Map Name (TypeBase dim o)
fs
supportsEquality (Scalar (Sum Map Name [TypeBase dim o]
fs)) = ([TypeBase dim o] -> Bool) -> Map Name [TypeBase dim o] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((TypeBase dim o -> Bool) -> [TypeBase dim o] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim o -> Bool
forall dim o. TypeBase dim o -> Bool
supportsEquality) Map Name [TypeBase dim o]
fs
supportsEquality (Scalar Arrow {}) = Bool
False

-- | Check that a type is non-functional, looking up the liftedness of type
-- variables.
orderZeroM :: [TypeParam] -> StructType -> TermTypeM Bool
orderZeroM :: [TypeParamBase VName] -> StructType -> TermTypeM Bool
orderZeroM [TypeParamBase VName]
tparams StructType
t = do
  (StructType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero StructType
t Bool -> Bool -> Bool
&&) (Bool -> Bool) -> ([Bool] -> Bool) -> [Bool] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> TermTypeM [Bool] -> TermTypeM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (QualName VName -> TermTypeM Bool)
-> [QualName VName] -> TermTypeM [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM QualName VName -> TermTypeM Bool
isUnlifted (StructType -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars StructType
t)
  where
    isUnlifted :: QualName VName -> TermTypeM Bool
isUnlifted QualName VName
qv = do
      case (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> Maybe (TypeParamBase VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qv) (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) [TypeParamBase VName]
tparams of
        Just (TypeParamType Liftedness
l VName
_ SrcLoc
_) -> Bool -> TermTypeM Bool
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> TermTypeM Bool) -> Bool -> TermTypeM Bool
forall a b. (a -> b) -> a -> b
$ Liftedness
l Liftedness -> Liftedness -> Bool
forall a. Ord a => a -> a -> Bool
< Liftedness
Lifted
        Maybe (TypeParamBase VName)
_ -> (Liftedness -> Liftedness -> Bool
forall a. Ord a => a -> a -> Bool
< Liftedness
Lifted) (Liftedness -> Bool) -> TermTypeM Liftedness -> TermTypeM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QualName VName -> TermTypeM Liftedness
lookupAbsTy QualName VName
qv

-- | Traverse the expression, emitting warnings and errors for various
-- problems:
--
-- * Unmatched cases.
--
-- * If any of the literals overflow their inferred types. Note:
--  currently unable to detect float underflow (such as 1e-400 -> 0)
--
-- * Function types appearing in places where they are not allowed (e.g.
--   returned from branches), and more generally lifted types used as
--   array elements.
--
-- The rationale is that it is easier to check for these things after all of the
-- type inference has been done, as they complicate the logic. Further, it is
-- also easier to produce good error messages here. The key is that we can only
-- enforce rules that do not affect type inference.
localChecks :: [TypeParam] -> Exp -> TermTypeM ()
localChecks :: [TypeParamBase VName] -> Exp -> TermTypeM ()
localChecks [TypeParamBase VName]
tparams Exp
orig_body = TermTypeM Exp -> TermTypeM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TermTypeM Exp -> TermTypeM ()) -> TermTypeM Exp -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Exp -> TermTypeM Exp
check Exp
orig_body
  where
    check :: Exp -> TermTypeM Exp
check e :: Exp
e@(AppExp (Match Exp
_ NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) (Info AppRes
rt)) = do
      ok <- [TypeParamBase VName] -> StructType -> TermTypeM Bool
orderZeroM [TypeParamBase VName]
tparams (AppRes -> StructType
appResType AppRes
rt)
      unless ok . typeError loc mempty $
        "Match-expression returns type"
          </> indent 2 (align (pretty (appResType rt)))
          </> "but match-results may not be of function type."
      let ps = (CaseBase Info VName -> PatBase Info VName StructType)
-> NonEmpty (CaseBase Info VName)
-> NonEmpty (PatBase Info VName StructType)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(CasePat PatBase Info VName StructType
p Exp
_ SrcLoc
_) -> PatBase Info VName StructType
p) NonEmpty (CaseBase Info VName)
cs
      case unmatched $ NE.toList ps of
        [] -> Exp -> TermTypeM Exp
recurse Exp
e
        [Match ()]
ps' ->
          SrcLoc -> Notes -> Doc () -> TermTypeM 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 SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM Exp)
-> (Doc () -> Doc ()) -> Doc () -> TermTypeM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"unmatched-cases" (Doc () -> TermTypeM Exp) -> Doc () -> TermTypeM Exp
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Unmatched cases in match expression:"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 ([Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
stack ((Match () -> Doc ()) -> [Match ()] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map Match () -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Match () -> Doc ann
pretty [Match ()]
ps'))
    check e :: Exp
e@(AppExp (If Exp
_ Exp
_ Exp
_ SrcLoc
loc) (Info AppRes
rt)) = do
      ok <- [TypeParamBase VName] -> StructType -> TermTypeM Bool
orderZeroM [TypeParamBase VName]
tparams (AppRes -> StructType
appResType AppRes
rt)
      unless ok . typeError loc mempty $
        "If-expression returns type"
          </> indent 2 (align (pretty (appResType rt)))
          </> "but if-results may not be of function type."
      recurse e
    check e :: Exp
e@(ArrayLit [Exp]
_ (Info StructType
t) SrcLoc
loc) = do
      (StructType -> TermTypeM ()) -> Maybe StructType -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (SrcLoc -> StructType -> TermTypeM ()
forall {o} {dim} {loc}.
(Pretty o, Pretty (Shape dim), Located loc) =>
loc -> TypeBase dim o -> TermTypeM ()
checkArrayElem SrcLoc
loc) (Maybe StructType -> TermTypeM ())
-> Maybe StructType -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Level -> StructType -> Maybe StructType
forall dim o. Level -> TypeBase dim o -> Maybe (TypeBase dim o)
peelArray Level
1 StructType
t
      Exp -> TermTypeM Exp
recurse Exp
e
    check e :: Exp
e@(AppExp (LetPat [SizeBinder VName]
_ PatBase Info VName StructType
p Exp
_ Exp
_ SrcLoc
_) Info AppRes
_) =
      PatBase Info VName StructType -> TermTypeM ()
forall (f :: * -> *).
MonadTypeChecker f =>
PatBase Info VName StructType -> f ()
mustBeIrrefutable PatBase Info VName StructType
p TermTypeM () -> TermTypeM Exp -> TermTypeM Exp
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Exp -> TermTypeM Exp
recurse Exp
e
    check e :: Exp
e@(AppExp (BinOp (QualName VName
v, SrcLoc
loc) Info StructType
_ (Exp
x, Info (Maybe VName)
_) (Exp, Info (Maybe VName))
_ SrcLoc
_) Info AppRes
_)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> VName
intrinsicVar Name
"==" = do
          case Exp -> StructType
typeOf Exp
x of
            Array {} -> do
              SrcLoc -> Doc () -> TermTypeM ()
forall loc. Located loc => loc -> Doc () -> TermTypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn SrcLoc
loc (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
                Text -> Doc ()
forall a. Text -> Doc a
textwrap
                  Text
"Comparing arrays with \"==\" is deprecated and will stop working in a future revision of the language."
            StructType
_ -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          SrcLoc -> StructType -> TermTypeM ()
forall {f :: * -> *} {loc} {o} {dim}.
(MonadTypeChecker f, Located loc, Pretty o, Pretty (Shape dim)) =>
loc -> TypeBase dim o -> f ()
checkEquality SrcLoc
loc (Exp -> StructType
typeOf Exp
x) TermTypeM () -> TermTypeM Exp -> TermTypeM Exp
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Exp -> TermTypeM Exp
recurse Exp
e
    check e :: Exp
e@(Var QualName VName
v (Info StructType
t) SrcLoc
loc)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> VName
intrinsicVar Name
"==" = do
          SrcLoc -> StructType -> TermTypeM ()
forall {f :: * -> *} {loc} {o} {dim}.
(MonadTypeChecker f, Located loc, Pretty o, Pretty (Shape dim)) =>
loc -> TypeBase dim o -> f ()
checkEquality SrcLoc
loc StructType
t TermTypeM () -> TermTypeM Exp -> TermTypeM Exp
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Exp -> TermTypeM Exp
recurse Exp
e
    check e :: Exp
e@(Lambda [Pat ParamType]
ps Exp
_ Maybe (TypeExp Exp VName)
_ Info ResRetType
_ SrcLoc
_) =
      (Pat ParamType -> TermTypeM ()) -> [Pat ParamType] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (PatBase Info VName StructType -> TermTypeM ()
forall (f :: * -> *).
MonadTypeChecker f =>
PatBase Info VName StructType -> f ()
mustBeIrrefutable (PatBase Info VName StructType -> TermTypeM ())
-> (Pat ParamType -> PatBase Info VName StructType)
-> Pat ParamType
-> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ParamType -> StructType)
-> Pat ParamType -> PatBase Info VName StructType
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct) [Pat ParamType]
ps TermTypeM () -> TermTypeM Exp -> TermTypeM Exp
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Exp -> TermTypeM Exp
recurse Exp
e
    check e :: Exp
e@(AppExp (LetFun (VName, SrcLoc)
_ ([TypeParamBase VName]
tparams', [Pat ParamType]
ps, Maybe (TypeExp Exp VName)
_, Info ResRetType
_, Exp
e1) Exp
e2 SrcLoc
_) Info AppRes
_) = do
      (Pat ParamType -> TermTypeM ()) -> [Pat ParamType] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (PatBase Info VName StructType -> TermTypeM ()
forall (f :: * -> *).
MonadTypeChecker f =>
PatBase Info VName StructType -> f ()
mustBeIrrefutable (PatBase Info VName StructType -> TermTypeM ())
-> (Pat ParamType -> PatBase Info VName StructType)
-> Pat ParamType
-> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ParamType -> StructType)
-> Pat ParamType -> PatBase Info VName StructType
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct) [Pat ParamType]
ps
      [TypeParamBase VName] -> Exp -> TermTypeM ()
localChecks ([TypeParamBase VName]
tparams' [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. Semigroup a => a -> a -> a
<> [TypeParamBase VName]
tparams) Exp
e1
      TermTypeM Exp -> TermTypeM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TermTypeM Exp -> TermTypeM ()) -> TermTypeM Exp -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Exp -> TermTypeM Exp
check Exp
e2
      Exp -> TermTypeM Exp
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
    check e :: Exp
e@(AppExp (Loop [VName]
_ Pat ParamType
p LoopInitBase Info VName
_ LoopFormBase Info VName
form Exp
_ SrcLoc
_) Info AppRes
_) = do
      PatBase Info VName StructType -> TermTypeM ()
forall (f :: * -> *).
MonadTypeChecker f =>
PatBase Info VName StructType -> f ()
mustBeIrrefutable ((ParamType -> StructType)
-> Pat ParamType -> PatBase Info VName StructType
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct Pat ParamType
p)
      case LoopFormBase Info VName
form of
        ForIn PatBase Info VName StructType
form_p Exp
_ -> PatBase Info VName StructType -> TermTypeM ()
forall (f :: * -> *).
MonadTypeChecker f =>
PatBase Info VName StructType -> f ()
mustBeIrrefutable PatBase Info VName StructType
form_p
        LoopFormBase Info VName
_ -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      ok <- [TypeParamBase VName] -> StructType -> TermTypeM Bool
orderZeroM [TypeParamBase VName]
tparams (Pat ParamType -> StructType
forall o. Pat (TypeBase Exp o) -> StructType
patternStructType Pat ParamType
p)
      unless ok . typeError (locOf p) mempty $
        "Loop parameter inferred to have type"
          </> indent 2 (align (pretty p))
          </> "but a loop parameter may not be of function type."
      recurse e
    check e :: Exp
e@(IntLit Integer
x Info StructType
ty SrcLoc
loc) =
      Exp
e Exp -> TermTypeM () -> TermTypeM Exp
forall a b. a -> TermTypeM b -> TermTypeM a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ case Info StructType
ty of
        Info (Scalar (Prim PrimType
t)) -> Bool -> Integer -> PrimType -> SrcLoc -> TermTypeM ()
forall {f :: * -> *} {loc} {a} {a}.
(MonadTypeChecker f, Located loc, Pretty a, Pretty a) =>
Bool -> a -> a -> loc -> f ()
errorBounds (Integer -> PrimType -> Bool
forall {a}. Integral a => a -> PrimType -> Bool
inBoundsI Integer
x PrimType
t) Integer
x PrimType
t SrcLoc
loc
        Info StructType
_ -> [Char] -> TermTypeM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Inferred type of int literal is not a number"
    check e :: Exp
e@(FloatLit Double
x Info StructType
ty SrcLoc
loc) =
      Exp
e Exp -> TermTypeM () -> TermTypeM Exp
forall a b. a -> TermTypeM b -> TermTypeM a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ case Info StructType
ty of
        Info (Scalar (Prim (FloatType FloatType
t))) -> Bool -> Double -> FloatType -> SrcLoc -> TermTypeM ()
forall {f :: * -> *} {loc} {a} {a}.
(MonadTypeChecker f, Located loc, Pretty a, Pretty a) =>
Bool -> a -> a -> loc -> f ()
errorBounds (Double -> FloatType -> Bool
forall {a}. RealFloat a => a -> FloatType -> Bool
inBoundsF Double
x FloatType
t) Double
x FloatType
t SrcLoc
loc
        Info StructType
_ -> [Char] -> TermTypeM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Inferred type of float literal is not a float"
    check e :: Exp
e@(Negate (IntLit Integer
x Info StructType
ty SrcLoc
loc1) SrcLoc
loc2) =
      Exp
e Exp -> TermTypeM () -> TermTypeM Exp
forall a b. a -> TermTypeM b -> TermTypeM a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ case Info StructType
ty of
        Info (Scalar (Prim PrimType
t)) -> Bool -> Integer -> PrimType -> SrcLoc -> TermTypeM ()
forall {f :: * -> *} {loc} {a} {a}.
(MonadTypeChecker f, Located loc, Pretty a, Pretty a) =>
Bool -> a -> a -> loc -> f ()
errorBounds (Integer -> PrimType -> Bool
forall {a}. Integral a => a -> PrimType -> Bool
inBoundsI (-Integer
x) PrimType
t) (-Integer
x) PrimType
t (SrcLoc
loc1 SrcLoc -> SrcLoc -> SrcLoc
forall a. Semigroup a => a -> a -> a
<> SrcLoc
loc2)
        Info StructType
_ -> [Char] -> TermTypeM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Inferred type of int literal is not a number"
    check Exp
e = Exp -> TermTypeM Exp
recurse Exp
e
    recurse :: Exp -> TermTypeM Exp
recurse = ASTMapper TermTypeM -> Exp -> TermTypeM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper TermTypeM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = check}

    checkEquality :: loc -> TypeBase dim o -> f ()
checkEquality loc
loc TypeBase dim o
t =
      Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (TypeBase dim o -> Bool
forall dim o. TypeBase dim o -> Bool
supportsEquality TypeBase dim o
t) (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$
        loc -> Notes -> Doc () -> f ()
forall loc a. Located loc => loc -> Notes -> Doc () -> f a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> f ()) -> Doc () -> f ()
forall a b. (a -> b) -> a -> b
$
          Doc ()
"Comparing equality of values of type"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 (TypeBase dim o -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase dim o -> Doc ann
pretty TypeBase dim o
t)
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"which does not support equality."

    -- Array elements must be unlifted: of non-varying size, and in
    -- particular not functions. This is a stricter requirement than
    -- 'orderZeroM', which permits size-lifted type parameters.
    checkArrayElem :: loc -> TypeBase dim o -> TermTypeM ()
checkArrayElem loc
loc TypeBase dim o
et = do
      Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (TypeBase dim o -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero TypeBase dim o
et) (TermTypeM () -> TermTypeM ())
-> (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. loc -> Notes -> Doc () -> TermTypeM ()
forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Type" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 (TypeBase dim o -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase dim o -> Doc ann
pretty TypeBase dim o
et) Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"found to be functional."
      (QualName VName -> TermTypeM ())
-> [QualName VName] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ QualName VName -> TermTypeM ()
checkElemVar ([QualName VName] -> TermTypeM ())
-> [QualName VName] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ TypeBase dim o -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars TypeBase dim o
et
      where
        checkElemVar :: QualName VName -> TermTypeM ()
checkElemVar QualName VName
qv = do
          l <- case (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> Maybe (TypeParamBase VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qv) (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) [TypeParamBase VName]
tparams of
            Just (TypeParamType Liftedness
l VName
_ SrcLoc
tploc) ->
              Either (Liftedness, Loc) Liftedness
-> TermTypeM (Either (Liftedness, Loc) Liftedness)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either (Liftedness, Loc) Liftedness
 -> TermTypeM (Either (Liftedness, Loc) Liftedness))
-> Either (Liftedness, Loc) Liftedness
-> TermTypeM (Either (Liftedness, Loc) Liftedness)
forall a b. (a -> b) -> a -> b
$ (Liftedness, Loc) -> Either (Liftedness, Loc) Liftedness
forall a b. a -> Either a b
Left (Liftedness
l, SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
tploc)
            Maybe (TypeParamBase VName)
_ -> Liftedness -> Either (Liftedness, Loc) Liftedness
forall a b. b -> Either a b
Right (Liftedness -> Either (Liftedness, Loc) Liftedness)
-> TermTypeM Liftedness
-> TermTypeM (Either (Liftedness, Loc) Liftedness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QualName VName -> TermTypeM Liftedness
lookupAbsTy QualName VName
qv
          case l of
            Left (Liftedness
l', Loc
tploc)
              | Liftedness
l' Liftedness -> Liftedness -> Bool
forall a. Eq a => a -> a -> Bool
/= Liftedness
Unlifted ->
                  loc -> Notes -> Doc () -> TermTypeM ()
forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
                    Doc ()
"Type parameter"
                      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (QualName VName -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName VName -> Doc ann
pretty QualName VName
qv)
                      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"bound at"
                      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> [Char] -> Doc ()
forall ann. [Char] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Loc -> [Char]
forall a. Located a => a -> [Char]
locStr Loc
tploc)
                      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"is lifted and cannot be an array element."
            Right Liftedness
l'
              | Liftedness
l' Liftedness -> Liftedness -> Bool
forall a. Eq a => a -> a -> Bool
/= Liftedness
Unlifted ->
                  loc -> Notes -> Doc () -> TermTypeM ()
forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
                    Doc ()
"Type"
                      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (QualName VName -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName VName -> Doc ann
pretty QualName VName
qv)
                      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"is lifted and cannot be an array element."
            Either (Liftedness, Loc) Liftedness
_ -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    bitWidth :: IntType -> Level
bitWidth IntType
ty = Level
8 Level -> Level -> Level
forall a. Num a => a -> a -> a
* IntType -> Level
forall a. Num a => IntType -> a
intByteSize IntType
ty :: Int

    inBoundsI :: a -> PrimType -> Bool
inBoundsI a
x (Signed IntType
t) = a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= -a
2 a -> Level -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (IntType -> Level
bitWidth IntType
t Level -> Level -> Level
forall a. Num a => a -> a -> a
- Level
1) Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2 a -> Level -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (IntType -> Level
bitWidth IntType
t Level -> Level -> Level
forall a. Num a => a -> a -> a
- Level
1)
    inBoundsI a
x (Unsigned IntType
t) = a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2 a -> Level -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ IntType -> Level
bitWidth IntType
t
    inBoundsI a
x (FloatType FloatType
Float16) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Half -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (a -> Half
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x :: Half)
    inBoundsI a
x (FloatType FloatType
Float32) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Float -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (a -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x :: Float)
    inBoundsI a
x (FloatType FloatType
Float64) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x :: Double)
    inBoundsI a
_ PrimType
Bool = [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"Inferred type of int literal is not a number"
    inBoundsF :: a -> FloatType -> Bool
inBoundsF a
x FloatType
Float16 = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Float -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (a -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x :: Float)
    inBoundsF a
x FloatType
Float32 = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Float -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (a -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x :: Float)
    inBoundsF a
x FloatType
Float64 = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
x

    errorBounds :: Bool -> a -> a -> loc -> f ()
errorBounds Bool
inBounds a
x a
ty loc
loc =
      Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
inBounds (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$
        loc -> Notes -> Doc () -> f ()
forall loc a. Located loc => loc -> Notes -> Doc () -> f a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> f ()) -> (Doc () -> Doc ()) -> Doc () -> f ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"literal-out-of-bounds" (Doc () -> f ()) -> Doc () -> f ()
forall a b. (a -> b) -> a -> b
$
          Doc ()
"Literal "
            Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> a -> Doc ()
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
x
            Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
" out of bounds for inferred type "
            Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> a -> Doc ()
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
ty
            Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."

-- | Check restrictions on recursive functions: the result must be first-order,
-- and any recursive applications must be invariant with respect to the
-- higher-order arguments. These are syntactic checks.
--
-- 'fname', 'params' and 'ret' describe the function whose body this is; if the
-- function is not recursive, this is a no-op.
recursionCheck ::
  [TypeParam] -> VName -> [Pat ParamType] -> ResType -> SrcLoc -> Exp -> TermTypeM ()
recursionCheck :: [TypeParamBase VName]
-> VName
-> [Pat ParamType]
-> ResType
-> SrcLoc
-> Exp
-> TermTypeM ()
recursionCheck [TypeParamBase VName]
tparams VName
fname [Pat ParamType]
params ResType
ret SrcLoc
fun_loc Exp
body =
  Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (VName
fname VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` FV -> Set VName
fvVars (Exp -> FV
freeInExp Exp
body)) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ do
    TermTypeM ()
checkRet
    higher_order <-
      (Pat ParamType -> TermTypeM Bool)
-> [Pat ParamType] -> TermTypeM [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Bool -> Bool) -> TermTypeM Bool -> TermTypeM Bool
forall a b. (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
not (TermTypeM Bool -> TermTypeM Bool)
-> (Pat ParamType -> TermTypeM Bool)
-> Pat ParamType
-> TermTypeM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeParamBase VName] -> StructType -> TermTypeM Bool
orderZeroM [TypeParamBase VName]
tparams (StructType -> TermTypeM Bool)
-> (Pat ParamType -> StructType) -> Pat ParamType -> TermTypeM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> StructType
forall o. Pat (TypeBase Exp o) -> StructType
patternStructType) [Pat ParamType]
params
    when (or higher_order) $ void $ check higher_order body
  where
    checkRet :: TermTypeM ()
checkRet = do
      ok <- [TypeParamBase VName] -> StructType -> TermTypeM Bool
orderZeroM [TypeParamBase VName]
tparams (StructType -> TermTypeM Bool) -> StructType -> TermTypeM Bool
forall a b. (a -> b) -> a -> b
$ ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
ret
      unless ok . typeError fun_loc mempty $
        "Recursive function"
          <+> dquotes (prettyName fname)
          <+> "returns type"
          </> indent 2 (pretty ret)
          </> "which is not first-order."
          </> "Write the function type as further parameters instead."

    check :: [Bool] -> Exp -> m Exp
check [Bool]
ho e :: Exp
e@(AppExp (Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
_) Info AppRes
_)
      | Var QualName VName
v Info StructType
_ SrcLoc
_ <- Exp
f,
        QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
fname = do
          [Bool] -> Loc -> [Exp] -> m ()
forall {m :: * -> *} {loc} {f :: * -> *}.
(MonadTypeChecker m, Located loc) =>
[Bool] -> loc -> [ExpBase f VName] -> m ()
checkRecApply [Bool]
ho (Exp -> Loc
forall a. Located a => a -> Loc
locOf Exp
e) ([Exp] -> m ()) -> [Exp] -> m ()
forall a b. (a -> b) -> a -> b
$ ((Info (Maybe VName), Exp) -> Exp)
-> [(Info (Maybe VName), Exp)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> Exp
forall a b. (a, b) -> b
snd ([(Info (Maybe VName), Exp)] -> [Exp])
-> [(Info (Maybe VName), Exp)] -> [Exp]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args
          ((Info (Maybe VName), Exp) -> m Exp)
-> NonEmpty (Info (Maybe VName), Exp) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Bool] -> Exp -> m Exp
check [Bool]
ho (Exp -> m Exp)
-> ((Info (Maybe VName), Exp) -> Exp)
-> (Info (Maybe VName), Exp)
-> m Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Info (Maybe VName), Exp) -> Exp
forall a b. (a, b) -> b
snd) NonEmpty (Info (Maybe VName), Exp)
args
          Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
    check [Bool]
_ (Var QualName VName
v Info StructType
_ SrcLoc
loc)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
fname =
          SrcLoc -> Notes -> Doc () -> m Exp
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 Exp) -> Doc () -> m Exp
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Recursive reference to"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
fname)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"must be a fully saturated application, as it has a higher-order parameter."
    check [Bool]
ho Exp
e = [Bool] -> Exp -> m Exp
recurse [Bool]
ho Exp
e
    recurse :: [Bool] -> Exp -> m Exp
recurse [Bool]
ho = ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper m
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = check ho}

    checkRecApply :: [Bool] -> loc -> [ExpBase f VName] -> m ()
checkRecApply [Bool]
ho loc
loc [ExpBase f VName]
args = do
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([ExpBase f VName] -> Level
forall a. [a] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [ExpBase f VName]
args Level -> Level -> Bool
forall a. Eq a => a -> a -> Bool
== [Pat ParamType] -> Level
forall a. [a] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [Pat ParamType]
params) (m () -> m ()) -> (Doc () -> m ()) -> Doc () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. loc -> Notes -> Doc () -> m ()
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Recursive application of"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
fname)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"is not fully saturated: expected"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Level -> Doc ()
forall ann. Level -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ([Pat ParamType] -> Level
forall a. [a] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [Pat ParamType]
params)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"arguments but got"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Level -> Doc ()
forall ann. Level -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ([ExpBase f VName] -> Level
forall a. [a] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [ExpBase f VName]
args)
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
      [(Bool, Pat ParamType, ExpBase f VName)]
-> ((Bool, Pat ParamType, ExpBase f VName) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Bool]
-> [Pat ParamType]
-> [ExpBase f VName]
-> [(Bool, Pat ParamType, ExpBase f VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Bool]
ho [Pat ParamType]
params [ExpBase f VName]
args) (((Bool, Pat ParamType, ExpBase f VName) -> m ()) -> m ())
-> ((Bool, Pat ParamType, ExpBase f VName) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Bool
is_ho, Pat ParamType
p, ExpBase f VName
arg) ->
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
is_ho (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
          case ExpBase f VName
arg of
            Var QualName VName
v f StructType
_ SrcLoc
_ | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames Pat ParamType
p -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            ExpBase f VName
_ ->
              Loc -> Notes -> Doc () -> m ()
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError (ExpBase f VName -> Loc
forall a. Located a => a -> Loc
locOf ExpBase f VName
arg) Notes
forall a. Monoid a => a
mempty (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
                Doc ()
"Higher-order argument in recursive application of"
                  Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
fname)
                  Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"must be passed unchanged, i.e. be the corresponding parameter."

-- | Instantiated sizes that unification has determined to be existential
-- ("pending"), and a mapping from pending copies to a representative: copies
-- from the same occurrence of an instantiated type parameter that were absorbed
-- from the same source denote the same existential size. See Note [Size
-- Inference].
pendingInstSizes :: Constraints -> (VName -> Bool, M.Map VName VName)
pendingInstSizes :: Constraints -> (VName -> Bool, Map VName VName)
pendingInstSizes Constraints
constraints = (VName -> Bool
pending, Map VName VName
reps)
  where
    key :: VName -> Maybe (Maybe VName)
key VName
v = case (Level, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((Level, Constraint) -> Constraint)
-> Maybe (Level, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Constraints -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Constraints
constraints of
      Just (ExistentialSize Maybe VName
k Maybe Loc
_ Usage
_) -> Maybe VName -> Maybe (Maybe VName)
forall a. a -> Maybe a
Just Maybe VName
k
      Maybe Constraint
_ -> Maybe (Maybe VName)
forall a. Maybe a
Nothing
    pending :: VName -> Bool
pending VName
v = case (Level, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((Level, Constraint) -> Constraint)
-> Maybe (Level, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Constraints -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Constraints
constraints of
      Just ExistentialSize {} -> Bool
True
      Just (CopySize VName
c Level
_ Usage
_) -> Maybe (Maybe VName) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Maybe VName) -> Bool) -> Maybe (Maybe VName) -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> Maybe (Maybe VName)
key VName
c
      Maybe Constraint
_ -> Bool
False
    groups :: Map (Level, VName) [VName]
groups =
      ([VName] -> [VName] -> [VName])
-> [((Level, VName), [VName])] -> Map (Level, VName) [VName]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith
        [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
(<>)
        [ ((Level
occ, VName
k), [VName
v])
        | (VName
v, (Level
_, CopySize VName
c Level
occ Usage
_)) <- Constraints -> [(VName, (Level, Constraint))]
forall k a. Map k a -> [(k, a)]
M.toList Constraints
constraints,
          Just (Just VName
k) <- [VName -> Maybe (Maybe VName)
key VName
c]
        ]
    reps :: Map VName VName
reps =
      [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
        [ (VName
v, VName
rep)
        | [VName]
vs <- Map (Level, VName) [VName] -> [[VName]]
forall k a. Map k a -> [a]
M.elems Map (Level, VName) [VName]
groups,
          let rep :: VName
rep = [VName] -> VName
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [VName]
vs,
          VName
v <- [VName]
vs,
          VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
rep
        ]

-- | Instantiated sizes (at or above the given level) that are still pending and
-- have not been determined to be existential behave like ordinary sizes from
-- here on: canonical sizes are plain size variables, and copies are equal to
-- their canonical variable. Existential ones are left alone; they are handled
-- by 'bindExistentialInsts'.
collapseInstSizes :: Level -> TermTypeM ()
collapseInstSizes :: Level -> TermTypeM ()
collapseInstSizes Level
min_lvl = do
  constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints
  let nonExistential VName
c = case (Level, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((Level, Constraint) -> Constraint)
-> Maybe (Level, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Constraints -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
c Constraints
constraints of
        Just ExistentialSize {} -> Bool
False
        Maybe Constraint
_ -> Bool
True
      collapse (Level
lvl, CopySize VName
c Level
_ Usage
usage)
        | Level
lvl Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
>= Level
min_lvl,
          VName -> Bool
nonExistential VName
c =
            (Level
lvl, Maybe Exp -> Usage -> Constraint
Size (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
c) (SrcLoc -> Exp) -> SrcLoc -> Exp
forall a b. (a -> b) -> a -> b
$ Usage -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Usage
usage) Usage
usage)
      collapse (Level
lvl, InstSize Liftedness
_ Usage
usage)
        | Level
lvl Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
>= Level
min_lvl = (Level
lvl, Maybe Exp -> Usage -> Constraint
Size Maybe Exp
forall a. Maybe a
Nothing Usage
usage)
      collapse (Level, Constraint)
x = (Level, Constraint)
x
  modifyConstraints $ M.map collapse

-- | Instantiated sizes that unification determined to be existential may remain
-- free in some types recorded in the AST - in particular the instantiated types
-- of higher-order functions, where the existential size occurs in the return
-- type of a function-typed parameter. Existentially bind such sizes at the
-- innermost possible position, mirroring what the type of the function argument
-- looks like.
bindExistentialInsts :: (ASTMappable e) => e -> TermTypeM e
bindExistentialInsts :: forall e. ASTMappable e => e -> TermTypeM e
bindExistentialInsts e
x = do
  -- 'pendingInstSizes' scans the entire constraint set, but its result
  -- is invariant across this traversal (any rigid sizes we introduce
  -- below are not instantiated sizes), so we compute it once instead of
  -- once per type in the AST.
  constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints
  let (pending, reps) = pendingInstSizes constraints
      repOf VName
v = Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> (VName -> Exp) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName SrcLoc
forall a. Monoid a => a
mempty (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> Subst t) -> Maybe VName -> Maybe (Subst t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
reps
      relevant VName
v = VName -> Bool
pending VName
v Bool -> Bool -> Bool
|| VName
v VName -> Map VName VName -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName VName
reps

      onType ::
        (Substitutable (TypeBase Size o)) =>
        TypeBase Size o ->
        TermTypeM (TypeBase Size o, [VName])
      onType TypeBase Exp o
t
        -- Fast path: this type mentions no pending or copied
        -- instantiated size, so 'applySubst'/'sizeFree' would be
        -- no-ops. Most types take this path.
        | Bool -> Bool
not ((VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any VName -> Bool
relevant (Set VName -> Bool) -> Set VName -> Bool
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> FV
forall o. TypeBase Exp o -> FV
freeInType TypeBase Exp o
t) = (TypeBase Exp o, [VName]) -> TermTypeM (TypeBase Exp o, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp o
t, [])
        | Bool
otherwise =
            SrcLoc
-> (Exp -> Maybe VName)
-> TypeBase Exp o
-> TermTypeM (TypeBase Exp o, [VName])
forall (m :: * -> *) o.
MonadUnify m =>
SrcLoc
-> (Exp -> Maybe VName)
-> TypeBase Exp o
-> m (TypeBase Exp o, [VName])
sizeFree SrcLoc
forall a. Monoid a => a
mempty ((VName -> Bool) -> Set VName -> Maybe VName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find VName -> Bool
pending (Set VName -> Maybe VName)
-> (Exp -> Set VName) -> Exp -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Set VName
fvVars (FV -> Set VName) -> (Exp -> FV) -> Exp -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> FV
freeInExp) (TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName]))
-> TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName])
forall a b. (a -> b) -> a -> b
$ TypeSubs -> TypeBase Exp o -> TypeBase Exp o
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
forall {t}. VName -> Maybe (Subst t)
repOf TypeBase Exp o
t

      onStruct ::
        (Substitutable (TypeBase Size o)) =>
        TypeBase Size o ->
        TermTypeM (TypeBase Size o)
      onStruct TypeBase Exp o
t = do
        (t', ext) <- TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName])
forall o.
Substitutable (TypeBase Exp o) =>
TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName])
onType TypeBase Exp o
t
        -- Existential sizes at the top level of a type have nowhere to
        -- be bound. Those that absorbed a rigid unknown size stand for
        -- a size that is actually computed at the recorded location,
        -- so they become rigid unknown sizes there, subjecting them to
        -- the causality check. The rest (absorbed from declared
        -- existentials, e.g. by a hole) are left alone.
        if null ext
          then pure t'
          else do
            let computedAt VName
v = case (Level, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((Level, Constraint) -> Constraint)
-> Maybe (Level, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Constraints -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Constraints
constraints of
                  Just (ExistentialSize Maybe VName
_ Maybe Loc
mloc Usage
_) -> Maybe Loc
mloc
                  Just (CopySize VName
c Level
_ Usage
_)
                    | Just (ExistentialSize Maybe VName
_ Maybe Loc
mloc Usage
_) <- (Level, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((Level, Constraint) -> Constraint)
-> Maybe (Level, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Constraints -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
c Constraints
constraints ->
                        Maybe Loc
mloc
                  Maybe Constraint
_ -> Maybe Loc
forall a. Maybe a
Nothing
            repls <- fmap (M.fromList . catMaybes) . forM (S.toList $ fvVars $ freeInType t) $ \VName
v ->
              case VName -> Maybe Loc
computedAt VName
v of
                Just Loc
dloc -> do
                  v' <- Loc -> RigidSource -> Name -> TermTypeM VName
forall a. Located a => a -> RigidSource -> Name -> TermTypeM VName
forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim Loc
dloc (Maybe (QualName VName) -> RigidSource
RigidRet Maybe (QualName VName)
forall a. Maybe a
Nothing) Name
"d"
                  pure $ Just (v, ExpSubst $ sizeFromName (qualName v') $ srclocOf dloc)
                Maybe Loc
Nothing -> Maybe (VName, Subst StructRetType)
-> TermTypeM (Maybe (VName, Subst StructRetType))
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (VName, Subst StructRetType)
forall a. Maybe a
Nothing
            pure $ applySubst (`M.lookup` repls) t

      tv =
        ASTMapper
          { mapOnExp :: Exp -> TermTypeM Exp
mapOnExp = ASTMapper TermTypeM -> Exp -> TermTypeM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper TermTypeM
tv,
            mapOnName :: QualName VName -> TermTypeM (QualName VName)
mapOnName = QualName VName -> TermTypeM (QualName VName)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnStructType :: StructType -> TermTypeM StructType
mapOnStructType = StructType -> TermTypeM StructType
forall o.
Substitutable (TypeBase Exp o) =>
TypeBase Exp o -> TermTypeM (TypeBase Exp o)
onStruct,
            mapOnParamType :: ParamType -> TermTypeM ParamType
mapOnParamType = ParamType -> TermTypeM ParamType
forall o.
Substitutable (TypeBase Exp o) =>
TypeBase Exp o -> TermTypeM (TypeBase Exp o)
onStruct,
            mapOnResRetType :: ResRetType -> TermTypeM ResRetType
mapOnResRetType = \(RetType [VName]
dims ResType
t) -> do
              (t', ext) <- ResType -> TermTypeM (ResType, [VName])
forall o.
Substitutable (TypeBase Exp o) =>
TypeBase Exp o -> TermTypeM (TypeBase Exp o, [VName])
onType ResType
t
              pure $ RetType (dims <> ext) t'
          }
  -- Global fast path: with no instantiated-size constraints at all, 'relevant'
  -- is false everywhere, so the traversal would just rebuild an identical copy.
  -- Skip it entirely - this is the common case.
  if any (isInstSize . snd) constraints
    then astMap tv x
    else pure x
  where
    isInstSize :: Constraint -> Bool
isInstSize ExistentialSize {} = Bool
True
    isInstSize CopySize {} = Bool
True
    isInstSize Constraint
_ = Bool
False

detectAmbiguousSizes :: TermTypeM ()
detectAmbiguousSizes :: TermTypeM ()
detectAmbiguousSizes = do
  Level -> TermTypeM ()
collapseInstSizes Level
0
  constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints
  mapM_ (notice constraints) $ M.toList constraints
  where
    -- Sizes that arise from instantiating inferred types have
    -- uninformative provenance. If a size variable with better
    -- provenance (e.g. a source-level size binder, or an
    -- instantiated size parameter) has been unified with the
    -- ambiguous size, we report that variable instead. This affects
    -- only the error message; which sizes are ambiguous is already
    -- settled.
    uninformative :: Usage -> Bool
uninformative (Usage Maybe Text
Nothing Loc
_) = Bool
True
    uninformative (Usage (Just Text
u) Loc
_) = Text
u Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text
"replaceTyVars", Text
"instantiation"]

    chase :: Map VName (a, Constraint) -> VName -> VName
chase Map VName (a, Constraint)
constraints VName
w = case (a, Constraint) -> Constraint
forall a b. (a, b) -> b
snd ((a, Constraint) -> Constraint)
-> Maybe (a, Constraint) -> Maybe Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (a, Constraint) -> Maybe (a, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
w Map VName (a, Constraint)
constraints of
      Just (Size (Just (Var QualName VName
w' Info StructType
_ SrcLoc
_)) Usage
_) -> Map VName (a, Constraint) -> VName -> VName
chase Map VName (a, Constraint)
constraints (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
w')
      Maybe Constraint
_ -> VName
w

    blame :: Map VName (a, Constraint) -> VName -> Usage -> (VName, Usage)
blame Map VName (a, Constraint)
constraints VName
v Usage
usage
      | Usage -> Bool
uninformative Usage
usage,
        (VName
v', Usage
usage') : [(VName, Usage)]
_ <-
          [ (VName
w, Usage
w_usage)
          | (VName
w, (a
_, Size (Just (Var QualName VName
w1 Info StructType
_ SrcLoc
_)) Usage
w_usage)) <- Map VName (a, Constraint) -> [(VName, (a, Constraint))]
forall k a. Map k a -> [(k, a)]
M.toList Map VName (a, Constraint)
constraints,
            Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Usage -> Bool
uninformative Usage
w_usage,
            Map VName (a, Constraint) -> VName -> VName
forall {a}. Map VName (a, Constraint) -> VName -> VName
chase Map VName (a, Constraint)
constraints (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
w1) VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v
          ] =
          (VName
v', Usage
usage')
      | Bool
otherwise = (VName
v, Usage
usage)

    notice :: Map VName (a, Constraint) -> (VName, (a, Constraint)) -> m ()
notice Map VName (a, Constraint)
constraints (VName
v, (a
_, Size Maybe Exp
Nothing Usage
usage)) =
      case Map VName (a, Constraint) -> VName -> Usage -> (VName, Usage)
forall {a}.
Map VName (a, Constraint) -> VName -> Usage -> (VName, Usage)
blame Map VName (a, Constraint)
constraints VName
v Usage
usage of
        (VName
v', Usage Maybe Text
Nothing Loc
loc) ->
          Loc -> Notes -> Doc () -> m ()
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> m ()) -> (Doc () -> Doc ()) -> Doc () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"ambiguous-size" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Ambiguous size" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v') Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
        (VName
v', Usage (Just Text
u) Loc
loc) ->
          Loc -> Notes -> Doc () -> m ()
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> m ()) -> (Doc () -> Doc ()) -> Doc () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"ambiguous-size" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Ambiguous size" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v') Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"arising from" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Text -> Doc ()
forall a. Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
u Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
    notice Map VName (a, Constraint)
_ (VName, (a, Constraint))
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | The names bound by these parameter patterns that are not the name of
-- a parameter itself, and hence cannot occur in the function type.
hiddenParamNames :: [Pat ParamType] -> [VName]
hiddenParamNames :: [Pat ParamType] -> [VName]
hiddenParamNames [Pat ParamType]
params = [VName]
hidden
  where
    param_all_names :: [VName]
param_all_names = [[VName]] -> [VName]
forall a. Monoid a => [a] -> a
mconcat ([[VName]] -> [VName]) -> [[VName]] -> [VName]
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> [VName]) -> [Pat ParamType] -> [[VName]]
forall a b. (a -> b) -> [a] -> [b]
map Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params
    named :: (PName, b, c) -> Maybe VName
named (Named VName
x, b
_, c
_) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
x
    named (PName
Unnamed, b
_, c
_) = Maybe VName
forall a. Maybe a
Nothing
    param_names :: Set VName
param_names =
      [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> Maybe VName) -> [Pat ParamType] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((PName, Diet, StructType) -> Maybe VName
forall {b} {c}. (PName, b, c) -> Maybe VName
named ((PName, Diet, StructType) -> Maybe VName)
-> (Pat ParamType -> (PName, Diet, StructType))
-> Pat ParamType
-> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> (PName, Diet, StructType)
forall d. Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam) [Pat ParamType]
params
    hidden :: [VName]
hidden = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Set VName -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Set VName
param_names) [VName]
param_all_names

-- | Rename the sizes bound by a type (parameter names and existential
-- quantifiers) to fresh names.
renameTypeBinders :: (Monoid o) => TypeBase Size o -> TermTypeM (TypeBase Size o)
renameTypeBinders :: forall o. Monoid o => TypeBase Exp o -> TermTypeM (TypeBase Exp o)
renameTypeBinders (Scalar (Arrow o
o PName
pn Diet
d StructType
pt (RetType [VName]
dims ResType
rt))) = do
  pt' <- StructType -> TermTypeM StructType
forall o. Monoid o => TypeBase Exp o -> TermTypeM (TypeBase Exp o)
renameTypeBinders StructType
pt
  (pn', pn_subst) <- case pn of
    Named VName
v -> do
      v' <- VName -> TermTypeM VName
forall (m :: * -> *). MonadTypeChecker m => VName -> m VName
newName VName
v
      pure (Named v', M.singleton v v')
    PName
Unnamed -> (PName, Map VName VName) -> TermTypeM (PName, Map VName VName)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PName
Unnamed, Map VName VName
forall a. Monoid a => a
mempty)
  dims' <- mapM newName dims
  let subst = Map VName VName
pn_subst Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
dims [VName]
dims')
      toSize VName
v = Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> Exp -> Subst t
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) SrcLoc
forall a. Monoid a => a
mempty
  rt' <- renameTypeBinders $ applySubst (fmap toSize . (`M.lookup` subst)) rt
  pure $ Scalar $ Arrow o pn' d pt' $ RetType dims' rt'
renameTypeBinders (Scalar (Record Map Name (TypeBase Exp o)
fs)) =
  ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> (Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o)
-> Map Name (TypeBase Exp o)
-> TypeBase Exp o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o) -> TypeBase Exp o)
-> TermTypeM (Map Name (TypeBase Exp o))
-> TermTypeM (TypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp o -> TermTypeM (TypeBase Exp o))
-> Map Name (TypeBase Exp o)
-> TermTypeM (Map Name (TypeBase Exp o))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse TypeBase Exp o -> TermTypeM (TypeBase Exp o)
forall o. Monoid o => TypeBase Exp o -> TermTypeM (TypeBase Exp o)
renameTypeBinders Map Name (TypeBase Exp o)
fs
renameTypeBinders (Scalar (Sum Map Name [TypeBase Exp o]
cs)) =
  ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> (Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o)
-> Map Name [TypeBase Exp o]
-> TypeBase Exp o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o] -> TypeBase Exp o)
-> TermTypeM (Map Name [TypeBase Exp o])
-> TermTypeM (TypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase Exp o] -> TermTypeM [TypeBase Exp o])
-> Map Name [TypeBase Exp o]
-> TermTypeM (Map Name [TypeBase Exp o])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse ((TypeBase Exp o -> TermTypeM (TypeBase Exp o))
-> [TypeBase Exp o] -> TermTypeM [TypeBase Exp o]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse TypeBase Exp o -> TermTypeM (TypeBase Exp o)
forall o. Monoid o => TypeBase Exp o -> TermTypeM (TypeBase Exp o)
renameTypeBinders) Map Name [TypeBase Exp o]
cs
renameTypeBinders TypeBase Exp o
t = TypeBase Exp o -> TermTypeM (TypeBase Exp o)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase Exp o
t

checkBinding ::
  ( VName,
    Maybe (TypeExp Exp VName),
    [TypeParam],
    [PatBase Info VName ParamType],
    ExpBase Info VName,
    SrcLoc
  ) ->
  TermTypeM
    ( [TypeParam],
      [Pat ParamType],
      Maybe (TypeExp Exp VName),
      ResRetType,
      Exp
    )
checkBinding :: (VName, Maybe (TypeExp Exp VName), [TypeParamBase VName],
 [Pat ParamType], Exp, SrcLoc)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
checkBinding (VName
fname, Maybe (TypeExp Exp VName)
maybe_retdecl, [TypeParamBase VName]
tparams, [Pat ParamType]
params, Exp
body, SrcLoc
loc) =
  TermTypeM
  ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
   ResRetType, Exp)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM
   ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
    ResRetType, Exp)
 -> TermTypeM
      ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
       ResRetType, Exp))
-> (([Pat ParamType]
     -> TermTypeM
          ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
           ResRetType, Exp))
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
          ResRetType, Exp))
-> ([Pat ParamType]
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
          ResRetType, Exp))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeParamBase VName]
-> [Pat ParamType]
-> ([Pat ParamType]
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
          ResRetType, Exp))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall a.
[TypeParamBase VName]
-> [Pat ParamType]
-> ([Pat ParamType] -> TermTypeM a)
-> TermTypeM a
bindingParams [TypeParamBase VName]
tparams [Pat ParamType]
params (([Pat ParamType]
  -> TermTypeM
       ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
        ResRetType, Exp))
 -> TermTypeM
      ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
       ResRetType, Exp))
-> ([Pat ParamType]
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
          ResRetType, Exp))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
    maybe_retdecl' <- (TypeExp Exp VName
 -> TermTypeM (TypeExp Exp VName, ResType, [VName]))
-> Maybe (TypeExp Exp VName)
-> TermTypeM (Maybe (TypeExp Exp VName, ResType, [VName]))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse TypeExp Exp VName
-> TermTypeM (TypeExp Exp VName, ResType, [VName])
checkTypeExpNonrigid Maybe (TypeExp Exp VName)
maybe_retdecl

    -- Bind the name in scope of its own body so it may recurse. Harmless even
    -- when the function is not actually recursive, as name resolution has
    -- hooked things up properly anyway. See Note [Checking recursive
    -- functions].
    self_binding <- case maybe_retdecl' of
      Just (TypeExp Exp VName
_, ResType
ret, [VName]
ext) ->
        [TypeParamBase VName] -> StructType -> ValBinding Exp
forall dim.
[TypeParamBase VName] -> TypeBase dim NoMode -> ValBinding dim
BoundV [TypeParamBase VName]
tparams (StructType -> ValBinding Exp)
-> TermTypeM StructType -> TermTypeM (ValBinding Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> TermTypeM StructType
forall o. Monoid o => TypeBase Exp o -> TermTypeM (TypeBase Exp o)
renameTypeBinders ([Pat ParamType] -> ResRetType -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
params' ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ext ResType
ret))
      Maybe (TypeExp Exp VName, ResType, [VName])
Nothing -> ValBinding Exp -> TermTypeM (ValBinding Exp)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ValBinding Exp
forall dim. ValBinding dim
RecursiveV
    (body', body_t) <-
      localScope (\TermScope Exp
scope -> TermScope Exp
scope {scopeVtable = M.insert fname self_binding $ scopeVtable scope}) $
        checkFunBody
          params'
          body
          ((\(TypeExp Exp VName
_, ResType
x, [VName]
_) -> ResType
x) <$> maybe_retdecl')
          (maybe loc srclocOf maybe_retdecl)

    params'' <- mapM updateTypes params'

    (maybe_retdecl'', rettype) <- case maybe_retdecl' of
      Just (TypeExp Exp VName
retdecl', ResType
ret, [VName]
_) -> do
        ret' <- ResType -> TermTypeM ResType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully ResType
ret
        pure (Just retdecl', ret')
      Maybe (TypeExp Exp VName, ResType, [VName])
Nothing ->
        (Maybe (TypeExp Exp VName), ResType)
-> TermTypeM (Maybe (TypeExp Exp VName), ResType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing, Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
body_t)

    verifyFunctionParams (Just fname) params''

    (tparams', params''', rettype') <-
      letGeneralise (baseName fname) loc tparams params''
        =<< unscopeUnknown rettype

    when
      ( null params
          && any isSizeParam tparams'
          && not (null (retDims rettype'))
      )
      $ typeError loc mempty
      $ textwrap "A size-polymorphic value binding may not have a type with an existential size."
        </> "Type of this binding is:"
        </> indent 2 (pretty rettype')
        </> "with the following type parameters:"
        </> indent 2 (sep $ map pretty $ filter isSizeParam tparams')

    pure (tparams', params''', maybe_retdecl'', rettype', body')

-- | Extract all the shape names that occur in positive position
-- (roughly, left side of an arrow) in a given type.
sizeNamesPos :: TypeBase Size o -> S.Set VName
sizeNamesPos :: forall o. TypeBase Exp o -> Set VName
sizeNamesPos (Scalar (Arrow o
_ PName
_ Diet
_ StructType
t1 (RetType [VName]
_ ResType
t2))) = StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
onParam StructType
t1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> ResType -> Set VName
forall o. TypeBase Exp o -> Set VName
sizeNamesPos ResType
t2
  where
    onParam :: TypeBase Size o -> S.Set VName
    onParam :: forall o. TypeBase Exp o -> Set VName
onParam (Scalar Arrow {}) = Set VName
forall a. Monoid a => a
mempty
    onParam (Scalar (Record Map Name (TypeBase Exp o)
fs)) = [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp o -> Set VName) -> [TypeBase Exp o] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Exp o -> Set VName
forall o. TypeBase Exp o -> Set VName
onParam ([TypeBase Exp o] -> [Set VName])
-> [TypeBase Exp o] -> [Set VName]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp o) -> [TypeBase Exp o]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase Exp o)
fs
    onParam (Scalar (TypeVar o
_ QualName VName
_ [TypeArg Exp]
targs)) = [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (TypeArg Exp -> Set VName) -> [TypeArg Exp] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg Exp -> Set VName
onTypeArg [TypeArg Exp]
targs
    onParam TypeBase Exp o
t = FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> FV
forall o. TypeBase Exp o -> FV
freeInType TypeBase Exp o
t
    onTypeArg :: TypeArg Exp -> Set VName
onTypeArg (TypeArgDim (Var QualName VName
d Info StructType
_ SrcLoc
_)) = VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
    onTypeArg (TypeArgDim Exp
_) = Set VName
forall a. Monoid a => a
mempty
    onTypeArg (TypeArgType StructType
t) = StructType -> Set VName
forall o. TypeBase Exp o -> Set VName
onParam StructType
t
sizeNamesPos TypeBase Exp o
_ = Set VName
forall a. Monoid a => a
mempty

-- | Verify certain restrictions on function parameters, and bail out
-- on dubious constructions.
--
-- These restrictions apply to all functions (anonymous or otherwise).
-- Top-level functions have further restrictions that are checked
-- during let-generalisation.
--
-- The parameters are assumed to already have their types normalised
-- ('updateTypes'), which both callers do immediately beforehand.
verifyFunctionParams :: Maybe VName -> [Pat ParamType] -> TermTypeM ()
verifyFunctionParams :: Maybe VName -> [Pat ParamType] -> TermTypeM ()
verifyFunctionParams Maybe VName
fname [Pat ParamType]
params =
  Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (Maybe Name -> Checking
CheckingParams (VName -> Name
baseName (VName -> Name) -> Maybe VName -> Maybe Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe VName
fname)) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    [VName] -> [Pat ParamType] -> TermTypeM ()
forall {m :: * -> *}.
MonadTypeChecker m =>
[VName] -> [Pat ParamType] -> m ()
verifyParams ((Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params) [Pat ParamType]
params
  where
    verifyParams :: [VName] -> [Pat ParamType] -> m ()
verifyParams [VName]
forbidden (Pat ParamType
p : [Pat ParamType]
ps)
      | VName
d : [VName]
_ <- (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
forbidden) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [VName]
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ Pat ParamType -> FV
forall o. Pat (TypeBase Exp o) -> FV
freeInPat Pat ParamType
p =
          Pat ParamType -> Notes -> Doc () -> m ()
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Pat ParamType
p Notes
forall a. Monoid a => a
mempty (Doc () -> m ()) -> (Doc () -> Doc ()) -> Doc () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"inaccessible-size" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Parameter"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Pat ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Pat ParamType -> Doc ann
pretty Pat ParamType
p)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"refers to size"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
d)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
forall ann. Doc ann
comma
                Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
"which will not be accessible to the caller"
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
forall ann. Doc ann
comma
                Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
"possibly because it is nested in a tuple or record."
                Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
"Consider ascribing an explicit type that does not reference "
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
d)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
      | Bool
otherwise = [VName] -> [Pat ParamType] -> m ()
verifyParams [VName]
forbidden' [Pat ParamType]
ps
      where
        forbidden' :: [VName]
forbidden' =
          case Pat ParamType -> (PName, Diet, StructType)
forall d. Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam Pat ParamType
p of
            (Named VName
v, Diet
_, StructType
_) -> VName -> [VName] -> [VName]
forall a. Eq a => a -> [a] -> [a]
delete VName
v [VName]
forbidden
            (PName, Diet, StructType)
_ -> [VName]
forbidden
    verifyParams [VName]
_ [] = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Move existentials down to the level where they are actually used
-- (i.e. have their "witnesses").  E.g. changes
--
-- @
-- ?[n].bool -> [n]bool
-- @
--
-- to
--
-- @
-- bool -> ?[n].[n]bool
-- @
injectExt :: [VName] -> TypeBase Size o -> RetTypeBase Size o
injectExt :: forall o. [VName] -> TypeBase Exp o -> RetTypeBase Exp o
injectExt [] TypeBase Exp o
ret = [VName] -> TypeBase Exp o -> RetTypeBase Exp o
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase Exp o
ret
injectExt [VName]
ext TypeBase Exp o
ret = [VName] -> TypeBase Exp o -> RetTypeBase Exp o
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ext_here (TypeBase Exp o -> RetTypeBase Exp o)
-> TypeBase Exp o -> RetTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> TypeBase Exp o
forall o. TypeBase Exp o -> TypeBase Exp o
deeper TypeBase Exp o
ret
  where
    (Set VName
immediate, Set VName
_) = TypeBase Exp o -> (Set VName, Set VName)
forall o. TypeBase Exp o -> (Set VName, Set VName)
dimUses TypeBase Exp o
ret
    ([VName]
ext_here, [VName]
ext_there) = (VName -> Bool) -> [VName] -> ([VName], [VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
immediate) [VName]
ext
    deeper :: TypeBase Size o -> TypeBase Size o
    deeper :: forall o. TypeBase Exp o -> TypeBase Exp o
deeper (Scalar (Prim PrimType
t)) = ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> ScalarTypeBase Exp o -> TypeBase Exp o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
t
    deeper (Scalar (Record Map Name (TypeBase Exp o)
fs)) = ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> ScalarTypeBase Exp o -> TypeBase Exp o
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o)
-> Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp o -> TypeBase Exp o)
-> Map Name (TypeBase Exp o) -> Map Name (TypeBase Exp o)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase Exp o -> TypeBase Exp o
forall o. TypeBase Exp o -> TypeBase Exp o
deeper Map Name (TypeBase Exp o)
fs
    deeper (Scalar (Sum Map Name [TypeBase Exp o]
cs)) = ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> ScalarTypeBase Exp o -> TypeBase Exp o
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o)
-> Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ ([TypeBase Exp o] -> [TypeBase Exp o])
-> Map Name [TypeBase Exp o] -> Map Name [TypeBase Exp o]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase Exp o -> TypeBase Exp o)
-> [TypeBase Exp o] -> [TypeBase Exp o]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Exp o -> TypeBase Exp o
forall o. TypeBase Exp o -> TypeBase Exp o
deeper) Map Name [TypeBase Exp o]
cs
    deeper (Scalar (Arrow o
als PName
p Diet
d1 StructType
t1 (RetType [VName]
t2_ext ResType
t2))) =
      ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> ScalarTypeBase Exp o -> TypeBase Exp o
forall a b. (a -> b) -> a -> b
$ o
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
als PName
p Diet
d1 StructType
t1 (ResRetType -> ScalarTypeBase Exp o)
-> ResRetType -> ScalarTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall o. [VName] -> TypeBase Exp o -> RetTypeBase Exp o
injectExt ([VName] -> [VName]
forall a. Ord a => [a] -> [a]
nubOrd ([VName]
ext_there [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
t2_ext)) ResType
t2
    deeper (Scalar (TypeVar o
o QualName VName
tn [TypeArg Exp]
targs)) =
      ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> ScalarTypeBase Exp o -> TypeBase Exp o
forall a b. (a -> b) -> a -> b
$ o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
tn ([TypeArg Exp] -> ScalarTypeBase Exp o)
-> [TypeArg Exp] -> ScalarTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ (TypeArg Exp -> TypeArg Exp) -> [TypeArg Exp] -> [TypeArg Exp]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg Exp -> TypeArg Exp
deeperArg [TypeArg Exp]
targs
    deeper t :: TypeBase Exp o
t@Array {} = TypeBase Exp o
t

    deeperArg :: TypeArg Exp -> TypeArg Exp
deeperArg (TypeArgType StructType
t) = StructType -> TypeArg Exp
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (StructType -> TypeArg Exp) -> StructType -> TypeArg Exp
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall o. TypeBase Exp o -> TypeBase Exp o
deeper StructType
t
    deeperArg (TypeArgDim Exp
d) = Exp -> TypeArg Exp
forall dim. dim -> TypeArg dim
TypeArgDim Exp
d

-- | Find all size variables in the given type that are covered by the
-- constraints, and produce size parameters that close over them.
--
-- The passed-in list of type parameters is always prepended to the
-- produced list of type parameters.
closeOverSizes ::
  Name ->
  SrcLoc ->
  [TypeParam] ->
  [StructType] ->
  ResType ->
  Constraints ->
  TermTypeM ([TypeParam], ResRetType)
closeOverSizes :: Name
-> SrcLoc
-> [TypeParamBase VName]
-> [StructType]
-> ResType
-> Constraints
-> TermTypeM ([TypeParamBase VName], ResRetType)
closeOverSizes Name
defname SrcLoc
defloc [TypeParamBase VName]
tparams [StructType]
paramts ResType
ret Constraints
substs = do
  (more_tparams, retext) <-
    [Either (TypeParamBase VName) VName]
-> ([TypeParamBase VName], [VName])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either (TypeParamBase VName) VName]
 -> ([TypeParamBase VName], [VName]))
-> ([Maybe (Either (TypeParamBase VName) VName)]
    -> [Either (TypeParamBase VName) VName])
-> [Maybe (Either (TypeParamBase VName) VName)]
-> ([TypeParamBase VName], [VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (Either (TypeParamBase VName) VName)]
-> [Either (TypeParamBase VName) VName]
forall a. [Maybe a] -> [a]
catMaybes
      ([Maybe (Either (TypeParamBase VName) VName)]
 -> ([TypeParamBase VName], [VName]))
-> TermTypeM [Maybe (Either (TypeParamBase VName) VName)]
-> TermTypeM ([TypeParamBase VName], [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((VName, Constraint)
 -> TermTypeM (Maybe (Either (TypeParamBase VName) VName)))
-> [(VName, Constraint)]
-> TermTypeM [Maybe (Either (TypeParamBase VName) VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (VName, Constraint)
-> TermTypeM (Maybe (Either (TypeParamBase VName) VName))
forall {f :: * -> *}.
(MonadUnify f, MonadTypeChecker f) =>
(VName, Constraint)
-> f (Maybe (Either (TypeParamBase VName) VName))
closeOver (Map VName Constraint -> [(VName, Constraint)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName Constraint -> [(VName, Constraint)])
-> Map VName Constraint -> [(VName, Constraint)]
forall a b. (a -> b) -> a -> b
$ ((Level, Constraint) -> Constraint)
-> Constraints -> Map VName Constraint
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Level, Constraint) -> Constraint
forall a b. (a, b) -> b
snd Constraints
to_close_over)
  let mkExt VName
v =
        case VName -> Constraints -> Maybe (Level, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Constraints
substs of
          Just (Level
_, UnknownSize {}) -> VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
          Maybe (Level, Constraint)
_ -> Maybe VName
forall a. Maybe a
Nothing

  pure
    ( tparams
        ++ more_tparams,
      injectExt (nubOrd $ retext ++ mapMaybe mkExt (S.toList $ fvVars $ freeInType ret)) ret
    )
  where
    -- Diet does not matter here.
    t :: StructType
t = [ParamType] -> ResRetType -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType ((StructType -> ParamType) -> [StructType] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe) [StructType]
paramts) (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] ResType
ret
    visible :: Set VName
visible = StructType -> Set VName
forall dim o. TypeBase dim o -> Set VName
typeVars StructType
t Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> FV -> Set VName
fvVars (StructType -> FV
forall o. TypeBase Exp o -> FV
freeInType StructType
t)
    to_close_over :: Constraints
to_close_over =
      (VName -> (Level, Constraint) -> Bool)
-> Constraints -> Constraints
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\VName
k (Level, Constraint)
_ -> VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
visible) Constraints
substs

    (Set VName
produced_sizes, Set VName
param_sizes) = StructType -> (Set VName, Set VName)
forall o. TypeBase Exp o -> (Set VName, Set VName)
dimUses StructType
t

    -- Avoid duplicate type parameters.
    closeOver :: (VName, Constraint)
-> f (Maybe (Either (TypeParamBase VName) VName))
closeOver (VName
k, Constraint
_)
      | VName
k VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (TypeParamBase VName -> VName) -> [TypeParamBase VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams =
          Maybe (Either (TypeParamBase VName) VName)
-> f (Maybe (Either (TypeParamBase VName) VName))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Either (TypeParamBase VName) VName)
forall a. Maybe a
Nothing
    closeOver (VName
k, Size Maybe Exp
Nothing Usage
_) =
      Maybe (Either (TypeParamBase VName) VName)
-> f (Maybe (Either (TypeParamBase VName) VName))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Either (TypeParamBase VName) VName)
 -> f (Maybe (Either (TypeParamBase VName) VName)))
-> Maybe (Either (TypeParamBase VName) VName)
-> f (Maybe (Either (TypeParamBase VName) VName))
forall a b. (a -> b) -> a -> b
$ Either (TypeParamBase VName) VName
-> Maybe (Either (TypeParamBase VName) VName)
forall a. a -> Maybe a
Just (Either (TypeParamBase VName) VName
 -> Maybe (Either (TypeParamBase VName) VName))
-> Either (TypeParamBase VName) VName
-> Maybe (Either (TypeParamBase VName) VName)
forall a b. (a -> b) -> a -> b
$ TypeParamBase VName -> Either (TypeParamBase VName) VName
forall a b. a -> Either a b
Left (TypeParamBase VName -> Either (TypeParamBase VName) VName)
-> TypeParamBase VName -> Either (TypeParamBase VName) VName
forall a b. (a -> b) -> a -> b
$ VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
TypeParamDim VName
k SrcLoc
forall a. Monoid a => a
mempty
    closeOver (VName
k, UnknownSize Loc
_ RigidSource
_)
      | VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
param_sizes,
        VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
produced_sizes = do
          notes <- SrcLoc -> Exp -> f Notes
forall a (m :: * -> *).
(Located a, MonadUnify m) =>
a -> Exp -> m Notes
dimNotes SrcLoc
defloc (Exp -> f Notes) -> Exp -> f Notes
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
k) SrcLoc
forall a. Monoid a => a
mempty
          typeError defloc notes . withIndexLink "unknown-param-def" $
            "Unknown size"
              <+> dquotes (prettyName k)
              <+> "in parameter of"
              <+> dquotes (prettyName defname)
              <> ", which is inferred as:"
                </> indent 2 (pretty t)
      | VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
produced_sizes =
          Maybe (Either (TypeParamBase VName) VName)
-> f (Maybe (Either (TypeParamBase VName) VName))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Either (TypeParamBase VName) VName)
 -> f (Maybe (Either (TypeParamBase VName) VName)))
-> Maybe (Either (TypeParamBase VName) VName)
-> f (Maybe (Either (TypeParamBase VName) VName))
forall a b. (a -> b) -> a -> b
$ Either (TypeParamBase VName) VName
-> Maybe (Either (TypeParamBase VName) VName)
forall a. a -> Maybe a
Just (Either (TypeParamBase VName) VName
 -> Maybe (Either (TypeParamBase VName) VName))
-> Either (TypeParamBase VName) VName
-> Maybe (Either (TypeParamBase VName) VName)
forall a b. (a -> b) -> a -> b
$ VName -> Either (TypeParamBase VName) VName
forall a b. b -> Either a b
Right VName
k
    closeOver (VName
_, Constraint
_) =
      Maybe (Either (TypeParamBase VName) VName)
-> f (Maybe (Either (TypeParamBase VName) VName))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Either (TypeParamBase VName) VName)
forall a. Maybe a
Nothing

letGeneralise ::
  Name ->
  SrcLoc ->
  [TypeParam] ->
  [Pat ParamType] ->
  ResType ->
  TermTypeM ([TypeParam], [Pat ParamType], ResRetType)
letGeneralise :: Name
-> SrcLoc
-> [TypeParamBase VName]
-> [Pat ParamType]
-> ResType
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
letGeneralise Name
defname SrcLoc
defloc [TypeParamBase VName]
tparams [Pat ParamType]
params ResType
restype =
  Checking
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (Name -> Checking
CheckingLetGeneralise Name
defname) (TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
 -> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType))
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
forall a b. (a -> b) -> a -> b
$ do
    cur_lvl <- TermTypeM Level
forall (m :: * -> *). MonadUnify m => m Level
curLevel
    collapseInstSizes $ cur_lvl - length params

    -- Re-normalise the types so that any instantiated sizes
    -- collapsed above are expressed in terms of their canonical
    -- variables, which can then be closed over.
    params' <- mapM updateTypes params
    restype' <- normTypeFully restype

    now_substs <- getConstraints

    -- Candidates for let-generalisation are those size variables that
    --
    -- (1) were not known before we checked this function, and
    --
    -- (2) are not used in the (new) definition of any size variables
    -- known before we checked this function.

    -- Criteria (1) and (2) is implemented by looking at the binding
    -- level of the size variables.
    let candidate (Level
lvl, b
_) = Level
lvl Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
>= (Level
cur_lvl Level -> Level -> Level
forall a. Num a => a -> a -> a
- [Pat ParamType] -> Level
forall a. [a] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [Pat ParamType]
params)
        new_substs = ((Level, Constraint) -> Bool) -> Constraints -> Constraints
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Level, Constraint) -> Bool
forall {b}. (Level, b) -> Bool
candidate Constraints
now_substs

    (tparams', RetType ret_dims restype'') <-
      closeOverSizes
        defname
        defloc
        tparams
        (map patternStructType params')
        restype'
        new_substs

    restype''' <- updateTypes restype''

    let used_sizes =
          ResType -> FV
forall o. TypeBase Exp o -> FV
freeInType ResType
restype''' FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (Pat ParamType -> FV) -> [Pat ParamType] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (ParamType -> FV
forall o. TypeBase Exp o -> FV
freeInType (ParamType -> FV)
-> (Pat ParamType -> ParamType) -> Pat ParamType -> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> ParamType
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType) [Pat ParamType]
params'
    case filter ((`S.notMember` fvVars used_sizes) . typeParamName) $
      filter isSizeParam tparams' of
      [] -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      TypeParamBase VName
tp : [TypeParamBase VName]
_ -> SizeBinder VName -> TermTypeM ()
forall (m :: * -> *) a.
MonadTypeChecker m =>
SizeBinder VName -> m a
unusedSize (SizeBinder VName -> TermTypeM ())
-> SizeBinder VName -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ VName -> SrcLoc -> SizeBinder VName
forall vn. vn -> SrcLoc -> SizeBinder vn
SizeBinder (TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase VName
tp) (TypeParamBase VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeParamBase VName
tp)

    -- We keep those type variables that were not closed over by
    -- let-generalisation.
    modifyConstraints $ M.filterWithKey $ \VName
k (Level, Constraint)
_ -> VName
k VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` (TypeParamBase VName -> VName) -> [TypeParamBase VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams'

    pure (tparams', params', RetType ret_dims restype''')

-- | Check the body of a function, and return it along with its type as
-- seen from outside the function: any 'hiddenParamNames' occurring as
-- sizes are replaced with fresh unknowns, which whoever decides the
-- function's return type then binds existentially.
checkFunBody ::
  [Pat ParamType] ->
  Exp ->
  Maybe ResType ->
  SrcLoc ->
  TermTypeM (Exp, StructType)
checkFunBody :: [Pat ParamType]
-> Exp -> Maybe ResType -> SrcLoc -> TermTypeM (Exp, StructType)
checkFunBody [Pat ParamType]
params Exp
body Maybe ResType
maybe_rettype SrcLoc
loc = do
  body' <- Exp -> TermTypeM Exp
checkExp Exp
body
  let hidden = [Pat ParamType] -> [VName]
hiddenParamNames [Pat ParamType]
params
  (body_t, _) <- unscopeType loc hidden =<< expTypeFully body'

  -- Unify body return type with return annotation, if one exists.
  case maybe_rettype of
    Just ResType
rettype ->
      case (VName -> Bool) -> Set VName -> Maybe VName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
hidden) (Set VName -> Maybe VName) -> Set VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ ResType -> FV
forall o. TypeBase Exp o -> FV
freeInType ResType
rettype of
        Just VName
v ->
          SrcLoc -> Notes -> Doc () -> TermTypeM ()
forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM 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 () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"The return type annotation"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Level -> Doc () -> Doc ()
forall ann. Level -> Doc ann -> Doc ann
indent Level
2 (Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
align (ResType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ResType -> Doc ann
pretty ResType
rettype))
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"refers to the name"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"which is bound to an inner component of a function parameter."
        Maybe VName
Nothing -> do
          let usage :: Usage
usage = Exp -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Exp
body Text
"return type annotation"
          Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (ResType -> StructType -> Checking
CheckingReturn ResType
rettype StructType
body_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 Usage
usage (ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
rettype) StructType
body_t
    Maybe ResType
Nothing -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  pure (body', body_t)

-- | Type-check a top-level (or module-level) function definition.
-- Despite the name, this is also used for checking constant
-- definitions, by treating them as 0-ary functions.
checkFunDef ::
  ( VName,
    Maybe (TypeExp (ExpBase NoInfo VName) VName),
    [TypeParam],
    [PatBase NoInfo VName ParamType],
    ExpBase NoInfo VName,
    SrcLoc
  ) ->
  TypeM
    ( [TypeParam],
      [Pat ParamType],
      Maybe (TypeExp Exp VName),
      ResRetType,
      Exp
    )
checkFunDef :: (VName, Maybe (TypeExp (ExpBase NoInfo VName) VName),
 [TypeParamBase VName], [PatBase NoInfo VName ParamType],
 ExpBase NoInfo VName, SrcLoc)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
checkFunDef (VName
fname, Maybe (TypeExp (ExpBase NoInfo VName) VName)
retdecl, [TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, ExpBase NoInfo VName
body, SrcLoc
loc) =
  (Either
   TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode)),
 [Pat ParamType], Maybe (TypeExp Exp VName), Exp)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
doChecks ((Either
    TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode)),
  [Pat ParamType], Maybe (TypeExp Exp VName), Exp)
 -> TypeM
      ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
       ResRetType, Exp))
-> TypeM
     (Either
        TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode)),
      [Pat ParamType], Maybe (TypeExp Exp VName), Exp)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName, Maybe (TypeExp (ExpBase NoInfo VName) VName),
 [TypeParamBase VName], [PatBase NoInfo VName ParamType],
 ExpBase NoInfo VName, SrcLoc)
-> TypeM
     (Either
        TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode)),
      [Pat ParamType], Maybe (TypeExp Exp VName), Exp)
Unsized.checkValDef (VName
fname, Maybe (TypeExp (ExpBase NoInfo VName) VName)
retdecl, [TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, ExpBase NoInfo VName
body, SrcLoc
loc)
  where
    doChecks :: (Either
   TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode)),
 [Pat ParamType], Maybe (TypeExp Exp VName), Exp)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
doChecks (Either
  TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode))
maybe_tysubsts, [Pat ParamType]
params', Maybe (TypeExp Exp VName)
retdecl', Exp
body') =
      case Either
  TypeError ([TypeParamBase VName], Map VName (TypeBase () NoMode))
maybe_tysubsts of
        Left TypeError
err -> TypeError
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall a. TypeError -> TypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError TypeError
err
        Right ([TypeParamBase VName]
generalised, Map VName (TypeBase () NoMode)
tysubsts) ->
          (Exp -> TermTypeM Exp)
-> Map VName (TypeBase () NoMode)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall a.
(Exp -> TermTypeM Exp)
-> Map VName (TypeBase () NoMode) -> TermTypeM a -> TypeM a
runTermTypeM Exp -> TermTypeM Exp
checkExp Map VName (TypeBase () NoMode)
tysubsts (TermTypeM
   ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
    ResRetType, Exp)
 -> TypeM
      ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
       ResRetType, Exp))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
forall a b. (a -> b) -> a -> b
$ do
            (tparams', params'', retdecl'', RetType dims rettype', body'') <-
              (VName, Maybe (TypeExp Exp VName), [TypeParamBase VName],
 [Pat ParamType], Exp, SrcLoc)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
      ResRetType, Exp)
checkBinding (VName
fname, Maybe (TypeExp Exp VName)
retdecl', [TypeParamBase VName]
generalised [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. Semigroup a => a -> a -> a
<> [TypeParamBase VName]
tparams, [Pat ParamType]
params', Exp
body', SrcLoc
loc)

            -- Since this is a top-level function, we also resolve overloaded
            -- types, using either defaults or complaining about ambiguities.
            detectAmbiguousSizes

            -- Then replace all inferred types in the body and parameters.
            body''' <- bindExistentialInsts =<< normTypeFully body''
            params''' <- mapM normTypeFully params''
            retdecl''' <- traverse updateTypes retdecl''
            rettype'' <- normTypeFully rettype'

            -- Check if the function body can actually be evaluated.
            causalityCheck body'''

            -- Check for various problems.
            mapM_ (mustBeIrrefutable . fmap toStruct) params''
            localChecks tparams' body'''
            recursionCheck tparams' fname params''' rettype'' loc body'''

            globals <- typeSchemes
            let ((body'''', updated_ret), errors) =
                  Consumption.checkValDef
                    globals
                    ( fname,
                      tparams',
                      params''',
                      body''',
                      RetType dims rettype'',
                      retdecl''',
                      loc
                    )

            mapM_ throwError errors

            pure (tparams', params''', retdecl''', updated_ret, body'''')

-- Note [Size Inference]
--
-- Type checking of terms is split into two passes. The unsized type checker
-- (Language.Futhark.TypeChecker.Terms.Unsized) infers types while ignoring
-- sizes entirely - its solution maps type variables to types whose dimensions
-- are all vacuous. The sized type checker (this module) receives that solution
-- (the 'termTyVars' field) and is responsible only for inferring sizes: the
-- concrete size of every dimension, and where existential quantifiers go. It
-- never re-infers anything besides sizes and where existential quantifiers go.
--
-- Whenever the size checker needs the type of something that the unsized
-- checker inferred, it instantiates the unsized type by replacing every
-- dimension with a fresh size variable ('instTyVars' when instantiating a type
-- scheme, 'replaceTyVars' elsewhere). Ordinary size unification then determines
-- what these variables stand for. This works out simply enough, except for
-- existential sizes.
--
-- ## Existential sizes
--
-- Consider
--
--   def (|>) '^a '^b (x: a) (f: a -> b) : b = f x
--
--   def main (xs: []i32) = xs |> filter (> 0)
--
-- where "b" is instantiated with the type of "filter (> 0)", which is [n]i32 ->
-- ?[m].[m]i32. At instantiation time we only know the unsized type []i32 for
-- "b" - existential sizes are invisible to the unsized pass, so we cannot know
-- that the size of "b" might be existential to the function. Worse, if the
-- instantiation does turn out to contain an existential size, then every
-- occurrence of "b" in the instantiated type scheme denotes a *distinct*
-- existential:
--
--   [n]i32 -> ([n]i32 -> ?[m].[m]i32) -> ?[m'].[m']i32
--
-- But if the instantiation turns out to have an ordinary size (say "xs |> map
-- (+1)"), all occurrences denote the *same* size, and we must not lose that
-- connection, or we would infer needlessly existential types.
--
-- We address this via three constraint forms (see 'Constraint'):
--
-- - InstSize: a canonical instantiated size, i.e. a dimension of the first
--   occurrence of an instantiated type parameter (or an instantiated size
--   parameter, which has the same nature).
--
-- - CopySize: a dimension of a later occurrence. These are given distinct names
--   precisely so that occurrences can become distinct existentials, but as long
--   as the size is not existential, a copy is equal to its canonical variable
--   (and unification treats it so, by redirecting links to the canonical
--   variable).
--
-- - ExistentialSize: When unification would otherwise fail by linking an
--   instantiated size to a size bound locally in the other type (an existential
--   or a parameter of a function type in the argument), it instead marks the
--   canonical variable with this ('unifySizes'). We say that the instantiated
--   size *absorbs* the locally bound size, and we call size variables that are
--   permitted to do so *absorbable* (see below). Absorption is refused for
--   unlifted type parameters, which cannot have existential sizes.
--
-- The pending existentials are then turned into proper sizes at the places that
-- can bind them:
--
-- - checkApply binds pending sizes in the return type of an application using
--   sizeFree: at the innermost RetType where possible, with the remainder
--   becoming existentials of the application itself (AppRes).
--
-- - bindExistentialInsts does the same for pending sizes that remain in types
--   recorded in the AST, in particular instantiated higher-order function
--   types, where the existential occurs in the return type of a function-typed
--   *parameter* and hence never passes through checkApply's return type.
--
-- - letGeneralise demotes instantiated sizes that are still pending and never
--   became existential to ordinary sizes (collapseInstSizes), at which point
--   they can be closed over as hidden size parameters. For local functions this
--   recreates the per-use size freshness that the old type checker obtained
--   from let-generalising type variables.
--
-- Two occurrences of an instantiated type parameter absorbed from the same
-- source must moreover denote the *same* existential, or we would infer
-- unwitnessed existentials where witnessed ones are possible. This is why
-- ExistentialSize records the size it was unified with and CopySize records
-- which occurrence it belongs to: pending sizes from the same occurrence with
-- the same source are given a single fresh name ('pendingInstSizes' in Terms).
--
-- ## Sizes bound by parameter patterns
--
-- An existential also arises whenever the type of a function body mentions a
-- name that a parameter *pattern* binds without it being the name of the
-- parameter itself, as the "k" of
--
--   \(k: i64, x: i32) -> replicate k x
--
-- where the parameter as a whole is anonymous. Such a name is not bound in the
-- function type, so we replace it with a fresh unknown size ('hiddenParamNames'
-- and 'unscopeType') in the function type. The logic is that we want to infer
-- the type that is visible from the "outside".
--
-- ## Dependent function types
--
-- A related problem is a dependent function type such as (n: i64) -> [n]i32 ->
-- [n]i32, whether reached by instantiating a type parameter with it or by
-- projecting it out of a value (e.g. a record field). The unsized pass preserves
-- parameter names, but the connection between the sizes and the binder is
-- exactly what was erased. Since every fresh size variable occurs exactly once,
-- and binders are cloned between occurrences of the type, it is safe to link a
-- fresh size to a binder of the type itself - such binders are registered
-- whenever the type is reconstructed from the erased solution ('registerBinders',
-- called from both 'instTyVars' and 'replaceTyVars'), and 'unifySizes' permits
-- exactly those links. Linking to any *other* locally bound size is what
-- signifies an existential (see above), or an error for sizes with no such
-- privileges. This binder linking is independent of absorption: it is available
-- even to the non-absorbable ('Unlifted') sizes of 'replaceTyVars', which is why
-- projecting a dependent function out of a record preserves its dependency
-- (tests/shapes/funshape11.fut).
--
-- ## Absorption privileges
--
-- To make the term of art explicit: a fresh size variable is *absorbable* if
-- unification may determine that it stands for an existential size. When an
-- absorbable size meets a size that is bound locally within the type it is
-- unified with, the mismatch is not an error; instead the absorbable size
-- absorbs the locally bound size - it is marked as a pending existential
-- (ExistentialSize), and is eventually existentially bound by the machinery
-- above. A size variable that is not absorbable must be resolved to an ordinary
-- size that is in scope, and encountering a locally bound size is an error for
-- it. Mechanically, absorbable sizes are exactly those constrained by InstSize
-- (or CopySize referring to one).
--
-- Which fresh sizes are absorbable is a fine line:
--
-- - Sizes arising from instantiation - of type parameters and of size
--   parameters - are absorbable. They occur exactly once, and stand for
--   "whatever size the context provides", which may well be existential.
--
-- - The sizes of a hole ('replaceTyVarsAbsorbable') are absorbable, as a hole
--   adopts whatever type the context provides. This includes adopting declared
--   existentials: a flexible existential that unification resolves to an
--   absorbable size is absorbed by it (see the end of the arrow case of
--   'unifyWith').
--
-- - The sizes of the inferred return type of a lambda and of a constructor
--   expression are likewise absorbable: a lambda returns whatever type the
--   context requires (in particular a lambda body ending in "#None" may well
--   have an existential option type), and the payload sizes of a constructor
--   application that are not determined by its arguments - notably those of
--   the *other* constructors - are adopted from the context. Relatedly, the
--   inferred return type of a lambda only hides (existentially binds) sizes
--   that are rigid; a still-flexible size has not been determined yet, and
--   hiding it would sever its connection to whatever the enclosing context
--   determines it to be (see 'inferReturnSizes').
--
-- - Sizes in the types of lambda parameters and patterns ('replaceTyVars') are
--   not absorbable (they are constrained 'InstSize' 'Unlifted', which can link
--   to binders of the type itself but not absorb an existential). They name the
--   sizes of actual bound values, and must be resolved to real sizes. For
--   example, this is what rejects
--
--     def f : (k: i64) -> [k]i32 -> i64 = \_ xs -> length xs
--
--   where the size of "xs" would otherwise silently absorb "k", which is not in
--   scope in the function body (tests/shapes/paramsize1.fut). Contrast with
--   tests/issue1168.fut, where the same shape of program must be accepted
--   because the inner function is size-generalised, and the *instantiation* of
--   its hidden size parameter is what absorbs the bound size of the expected
--   type.
--
-- ## Absorption and causality
--
-- Absorbing an existential is not always innocent. Whether the size that was
-- absorbed is a *declared* existential (of an ascribed type) or a *rigid
-- unknown* size (one whose value is computed at a specific point in the
-- program, such as the result of applying a function with an existential
-- return type) makes a difference:
--
--   def ite b t f = if b then t () else f ()
--   def f : () -> option ([]i32) = \() -> #None                    -- fine
--   def g b = ite b (\() -> #None) (\() -> #Some (gen ()))         -- rejected
--
-- In g, the type of "#None" is forced (via the instantiation of the type
-- parameter of "ite") to have the size produced by "gen ()" inside the other
-- lambda - a size that is only computed elsewhere, so the constructed value
-- cannot know its payload size (tests/sumtypes/sumtype52.fut). ExistentialSize
-- therefore records the location when the absorbed size was rigid, and when
-- such a pending existential remains in a type with no position to bind it,
-- 'bindExistentialInsts' turns it into a rigid unknown size at the recorded
-- location. The causality check then rejects expressions that need it (a sum
-- constructor, say) before it is computed, with no knowledge of this
-- machinery. Two subtleties: the ext variables temporarily introduced in the
-- arrow case of 'unifyWith' shadow the registered constraints of the binders,
-- so the rigidity of a binder is determined from the constraints as they were
-- before ('rigidPre'); and a pending size that checkApply has already bound in
-- a remaining parameter type (of a partially applied function) is registered
-- as a rigid unknown size by 'sizeFree', which is what carries the obligation
-- across applications.

-- Note [Checking recursive functions]
--
-- A function may refer to itself in its own body. The difficulty is that the
-- function's type is not fully known until we have checked that body, so we
-- cannot simply look the name up like any other. We follow the textbook
-- Hindley-Milner treatment on monomoprhic recursion, where all occurrences of
-- the function within its own body share a single type. However, we add the
-- twist that *size parameters* are allowed to differ.
--
-- The handling is spread across name resolution and the two type-checking
-- passes:
--
-- 1. Name resolution brings the function name into scope of its own body, but
--    only for *syntactic* functions (those with parameters).
--
-- 2. The unsized pass binds the name to a fresh monomorphic type variable while
--    checking the body, then emits a single constraint equating that variable
--    with the actual function type. This is standard monomorphic recursion.
--
-- 3. The sized pass considers two cases, distinguished by whether the function
--    has a declared return type:
--
--    * With a declared return type, the function's size-precise type is known
--      before we check the body, so we bind the name to that type scheme
--      ('BoundV'). Recursive occurrences are then instantiated like calls to
--      any other function: sizes are refreshed per occurrence, but the size
--      *relationships* of the signature are kept - e.g. that '[n]i32 -> [n]i32'
--      returns an array the size of its argument. This permits size-polymorphic
--      recursion.
--
--    * Without a declared return type, we cannot know the return type in
--      advance, so we bind the name to 'RecursiveV' . 'lookupVar' then resolves
--      each occurrence with 'replaceTyVars' on the type the unsized pass
--      recorded, which creates *fresh, unrelated* sizes for every dimension,
--      which is OK whenever no size relationship in the return type matters.
--      This means we cannot *infer* size constraints for recursive functions.
--
-- Only top-level self-recursion is handled. Mutual recursion and local
-- (let-bound) recursion are not.