-- | Facilities for type-checking Futhark terms.  Checking a term
-- requires a little more context to track uniqueness 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 uniqueness
-- types.  This is mostly done in an ad hoc way, and many programs
-- will require the programmer to fall back on type annotations.
module Language.Futhark.TypeChecker.Terms
  ( checkOneExp,
    checkSizeExp,
    checkFunDef,
  )
where

import Control.Monad
import Control.Monad.Except
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.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.Util (mapAccumLM, nubOrd)
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, lookupMod)
import Language.Futhark.TypeChecker.Terms.Loop
import Language.Futhark.TypeChecker.Terms.Monad
import Language.Futhark.TypeChecker.Terms.Pat
import Language.Futhark.TypeChecker.Types
import Language.Futhark.TypeChecker.Unify
import Prelude hiding (mod)

hasBinding :: Exp -> Bool
hasBinding :: Size -> 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 Size
e = Maybe Size -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe Size -> Bool) -> Maybe Size -> Bool
forall a b. (a -> b) -> a -> b
$ ASTMapper Maybe -> Size -> Maybe Size
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Size -> m Size
astMap ASTMapper Maybe
m Size
e
  where
    m :: ASTMapper Maybe
m =
      ASTMapper Maybe
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = \Size
e' -> if Size -> Bool
hasBinding Size
e' then Maybe Size
forall a. Maybe a
Nothing else Size -> Maybe Size
forall a. a -> Maybe a
Just Size
e'}

overloadedTypeVars :: Constraints -> Names
overloadedTypeVars :: Constraints -> Set VName
overloadedTypeVars = [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName)
-> (Constraints -> [Set VName]) -> Constraints -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Constraint) -> Set VName)
-> [(Int, Constraint)] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Constraint) -> Set VName
forall {a}. (a, Constraint) -> Set VName
f ([(Int, Constraint)] -> [Set VName])
-> (Constraints -> [(Int, Constraint)])
-> Constraints
-> [Set VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constraints -> [(Int, Constraint)]
forall k a. Map k a -> [a]
M.elems
  where
    f :: (a, Constraint) -> Set VName
f (a
_, HasFields Liftedness
_ Map Name StructType
fs Usage
_) = [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
$ (StructType -> Set VName) -> [StructType] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars ([StructType] -> [Set VName]) -> [StructType] -> [Set VName]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [StructType]
forall k a. Map k a -> [a]
M.elems Map Name StructType
fs
    f (a, Constraint)
_ = Set VName
forall a. Monoid a => a
mempty

--- Basic checking

-- | Determine if the two types are identical, ignoring uniqueness.
-- 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 -> Size -> Size -> TermTypeM (StructType, [VName])
unifyBranches SrcLoc
loc Size
e1 Size
e2 = do
  StructType
e1_t <- Size -> TermTypeM StructType
expTypeFully Size
e1
  StructType
e2_t <- Size -> TermTypeM StructType
expTypeFully Size
e2
  SrcLoc
-> StructType -> StructType -> TermTypeM (StructType, [VName])
unifyBranchTypes SrcLoc
loc StructType
e1_t StructType
e2_t

sliceShape ::
  Maybe (SrcLoc, Rigidity) ->
  [DimIndex] ->
  TypeBase Size as ->
  TermTypeM (TypeBase Size as, [VName])
sliceShape :: forall as.
Maybe (SrcLoc, Rigidity)
-> [DimIndex]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
sliceShape Maybe (SrcLoc, Rigidity)
r [DimIndex]
slice t :: TypeBase Size as
t@(Array as
u (Shape [Size]
orig_dims) ScalarTypeBase Size NoUniqueness
et) =
  StateT [VName] TermTypeM (TypeBase Size as)
-> [VName] -> TermTypeM (TypeBase Size as, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ([Size] -> TypeBase Size as
setDims ([Size] -> TypeBase Size as)
-> StateT [VName] TermTypeM [Size]
-> StateT [VName] TermTypeM (TypeBase Size as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DimIndex] -> [Size] -> StateT [VName] TermTypeM [Size]
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
[DimIndex] -> [Size] -> t TermTypeM [Size]
adjustDims [DimIndex]
slice [Size]
orig_dims) []
  where
    setDims :: [Size] -> TypeBase Size as
setDims [] = Int -> TypeBase Size as -> TypeBase Size as
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray ([Size] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Size]
orig_dims) TypeBase Size as
t
    setDims [Size]
dims' = as
-> Shape Size
-> ScalarTypeBase Size NoUniqueness
-> TypeBase Size as
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array as
u ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [Size]
dims') ScalarTypeBase Size NoUniqueness
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 :: Size -> Maybe Size -> Maybe Size -> Maybe Size -> t TermTypeM Size
sliceSize Size
orig_d Maybe Size
i Maybe Size
j Maybe Size
stride =
      case Maybe (SrcLoc, Rigidity)
r of
        Just (SrcLoc
loc, Rigid RigidSource
_) -> do
          (Size
d, Maybe VName
ext) <-
            TermTypeM (Size, Maybe VName) -> t TermTypeM (Size, 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 (Size, Maybe VName) -> t TermTypeM (Size, Maybe VName))
-> (SizeSource -> TermTypeM (Size, Maybe VName))
-> SizeSource
-> t TermTypeM (Size, Maybe VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> SizeSource -> TermTypeM (Size, Maybe VName)
extSize SrcLoc
loc (SizeSource -> t TermTypeM (Size, Maybe VName))
-> SizeSource -> t TermTypeM (Size, Maybe VName)
forall a b. (a -> b) -> a -> b
$
              Maybe Size
-> Maybe (ExpBase NoInfo VName)
-> Maybe (ExpBase NoInfo VName)
-> Maybe (ExpBase NoInfo VName)
-> SizeSource
SourceSlice Maybe Size
orig_d' (Size -> ExpBase NoInfo VName
bareExp (Size -> ExpBase NoInfo VName)
-> Maybe Size -> Maybe (ExpBase NoInfo VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Size
i) (Size -> ExpBase NoInfo VName
bareExp (Size -> ExpBase NoInfo VName)
-> Maybe Size -> Maybe (ExpBase NoInfo VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Size
j) (Size -> ExpBase NoInfo VName
bareExp (Size -> ExpBase NoInfo VName)
-> Maybe Size -> Maybe (ExpBase NoInfo VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Size
stride)
          ([VName] -> [VName]) -> t TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Maybe VName -> [VName]
forall a. Maybe a -> [a]
maybeToList Maybe VName
ext ++)
          Size -> t TermTypeM Size
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d
        Just (SrcLoc
loc, Rigidity
Nonrigid) ->
          TermTypeM Size -> t TermTypeM Size
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 Size -> t TermTypeM Size)
-> TermTypeM Size -> t TermTypeM Size
forall a b. (a -> b) -> a -> b
$
            (QualName VName -> SrcLoc -> Size)
-> SrcLoc -> QualName VName -> Size
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Size
sizeFromName SrcLoc
loc (QualName VName -> Size)
-> (VName -> QualName VName) -> VName -> Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName
              (VName -> Size) -> TermTypeM VName -> TermTypeM Size
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
          VName
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"
          ([VName] -> [VName]) -> t TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName
v :)
          Size -> t TermTypeM Size
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> t TermTypeM Size) -> Size -> t TermTypeM Size
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Size
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) SrcLoc
forall a. Monoid a => a
mempty
      where
        -- The original size does not matter if the slice is fully specified.
        orig_d' :: Maybe Size
orig_d'
          | Maybe Size -> Bool
forall a. Maybe a -> Bool
isJust Maybe Size
i, Maybe Size -> Bool
forall a. Maybe a -> Bool
isJust Maybe Size
j = Maybe Size
forall a. Maybe a
Nothing
          | Bool
otherwise = Size -> Maybe Size
forall a. a -> Maybe a
Just Size
orig_d

    warnIfBinding :: Bool
-> Size
-> Maybe Size
-> Maybe Size
-> Maybe Size
-> Size
-> t TermTypeM ([Size] -> [Size])
warnIfBinding Bool
binds Size
d Maybe Size
i Maybe Size
j Maybe Size
stride Size
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 (Size -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Size
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."
          (:) (Size -> [Size] -> [Size])
-> t TermTypeM Size -> t TermTypeM ([Size] -> [Size])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Size -> Maybe Size -> Maybe Size -> Maybe Size -> t TermTypeM Size
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t, MonadState [VName] (t TermTypeM)) =>
Size -> Maybe Size -> Maybe Size -> Maybe Size -> t TermTypeM Size
sliceSize Size
d Maybe Size
i Maybe Size
j Maybe Size
stride
        else ([Size] -> [Size]) -> t TermTypeM ([Size] -> [Size])
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size
size :)

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

    sizeMinus :: Size -> Size -> Size
sizeMinus Size
j Size
i =
      AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        ( (QualName VName, SrcLoc)
-> Info StructType
-> (Size, Info (Maybe VName))
-> (Size, 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
sizeBinOpInfo
            (Size
j, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
            (Size
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 -> Size) -> Info AppRes -> Size
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} {u}. TypeBase dim u
i64 []
    i64 :: TypeBase dim u
i64 = ScalarTypeBase dim u -> TypeBase dim u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim u -> TypeBase dim u)
-> ScalarTypeBase dim u -> TypeBase dim u
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase dim u)
-> PrimType -> ScalarTypeBase dim u
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    sizeBinOpInfo :: Info StructType
sizeBinOpInfo = StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ [ParamType] -> ResRetType -> StructType
foldFunType [ParamType
forall {dim} {u}. TypeBase dim u
i64, ParamType
forall {dim} {u}. TypeBase dim u
i64] (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Size Uniqueness
forall {dim} {u}. TypeBase dim u
i64
sliceShape Maybe (SrcLoc, Rigidity)
_ [DimIndex]
_ TypeBase Size as
t = (TypeBase Size as, [VName])
-> TermTypeM (TypeBase Size as, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Size as
t, [])

--- Main checkers

checkAscript ::
  SrcLoc ->
  TypeExp (ExpBase NoInfo VName) VName ->
  ExpBase NoInfo VName ->
  TermTypeM (TypeExp Exp VName, Exp)
checkAscript :: SrcLoc
-> TypeExp (ExpBase NoInfo VName) VName
-> ExpBase NoInfo VName
-> TermTypeM (TypeExp Size VName, Size)
checkAscript SrcLoc
loc TypeExp (ExpBase NoInfo VName) VName
te ExpBase NoInfo VName
e = do
  (TypeExp Size VName
te', TypeBase Size Uniqueness
decl_t, [VName]
_) <- TypeExp (ExpBase NoInfo VName) VName
-> TermTypeM
     (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
checkTypeExpNonrigid TypeExp (ExpBase NoInfo VName) VName
te
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  StructType
e_t <- Size -> TermTypeM StructType
expTypeFully Size
e'

  Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (StructType -> StructType -> Checking
CheckingAscription (TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
decl_t) StructType
e_t) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"type ascription") (TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
decl_t) StructType
e_t

  (TypeExp Size VName, Size) -> TermTypeM (TypeExp Size VName, Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp Size VName
te', Size
e')

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

  StructType
te_t_nonrigid <- [VName] -> StructType -> TermTypeM StructType
forall {t :: * -> * -> *} {t :: * -> *} {d}.
(Bitraversable t, Foldable t) =>
t VName -> t Size d -> TermTypeM (t Size d)
makeNonExtFresh [VName]
ext (StructType -> TermTypeM StructType)
-> StructType -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
te_t

  Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (StructType -> StructType -> Checking
CheckingAscription (TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
te_t) StructType
e_t) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"size coercion") StructType
e_t StructType
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.
  (TypeExp Size VName, StructType, Size)
-> TermTypeM (TypeExp Size VName, StructType, Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp Size VName
te', TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
te_t, Size
e')
  where
    makeNonExtFresh :: t VName -> t Size d -> TermTypeM (t Size d)
makeNonExtFresh t VName
ext = (Size -> TermTypeM Size)
-> (d -> TermTypeM d) -> t Size d -> TermTypeM (t Size 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 Size -> TermTypeM Size
onDim d -> TermTypeM d
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      where
        onDim :: Size -> TermTypeM Size
onDim d :: Size
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 = Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d
        onDim Size
d = do
          VName
v <- Name -> TermTypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
"coerce"
          VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ())
-> (Usage -> Constraint) -> Usage -> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Size -> Usage -> Constraint
Size Maybe Size
forall a. Maybe a
Nothing (Usage -> TermTypeM ()) -> Usage -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
            SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage
              SrcLoc
loc
              Text
"a size coercion where the underlying expression size cannot be determined"
          Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Size
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) (Size -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Size
d)

-- Used to 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 u ->
  TermTypeM (TypeBase Size u)
unscopeUnknown :: forall u. TypeBase Size u -> TermTypeM (TypeBase Size u)
unscopeUnknown TypeBase Size u
t = do
  Constraints
constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints
  -- These sizes will be immediately turned into existentials, so we
  -- do not need to care about their location.
  (TypeBase Size u, [VName]) -> TypeBase Size u
forall a b. (a, b) -> a
fst ((TypeBase Size u, [VName]) -> TypeBase Size u)
-> TermTypeM (TypeBase Size u, [VName])
-> TermTypeM (TypeBase Size u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc
-> (Size -> Maybe VName)
-> TypeBase Size u
-> TermTypeM (TypeBase Size u, [VName])
forall (m :: * -> *) u.
MonadUnify m =>
SrcLoc
-> (Size -> Maybe VName)
-> TypeBase Size u
-> m (TypeBase Size u, [VName])
sizeFree SrcLoc
forall a. Monoid a => a
mempty (Constraints -> Size -> Maybe VName
forall {a}. Map VName (a, Constraint) -> Size -> Maybe VName
expKiller Constraints
constraints) TypeBase Size u
t
  where
    expKiller :: Map VName (a, Constraint) -> Size -> Maybe VName
expKiller Map VName (a, Constraint)
_ Var {} = Maybe VName
forall a. Maybe a
Nothing
    expKiller Map VName (a, Constraint)
constraints Size
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
$ Size -> FV
freeInExp Size
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 Size u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size u
t

unscopeType ::
  SrcLoc ->
  [VName] ->
  TypeBase Size as ->
  TermTypeM (TypeBase Size as, [VName])
unscopeType :: forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType SrcLoc
tloc [VName]
unscoped =
  SrcLoc
-> (Size -> Maybe VName)
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
forall (m :: * -> *) u.
MonadUnify m =>
SrcLoc
-> (Size -> Maybe VName)
-> TypeBase Size u
-> m (TypeBase Size u, [VName])
sizeFree SrcLoc
tloc ((Size -> Maybe VName)
 -> TypeBase Size as -> TermTypeM (TypeBase Size as, [VName]))
-> (Size -> Maybe VName)
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
forall a b. (a -> b) -> a -> b
$ (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)
-> (Size -> Set VName) -> Size -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Set VName
fvVars (FV -> Set VName) -> (Size -> FV) -> Size -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Size -> FV
freeInExp

checkExp :: ExpBase NoInfo VName -> TermTypeM Exp
checkExp :: ExpBase NoInfo VName -> TermTypeM Size
checkExp (Literal PrimValue
val SrcLoc
loc) =
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> Size
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal PrimValue
val SrcLoc
loc
checkExp (Hole NoInfo StructType
_ SrcLoc
loc) = do
  StructType
t <- SrcLoc -> Name -> TermTypeM StructType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc
checkExp (StringLit [Word8]
vs SrcLoc
loc) =
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [Word8] -> SrcLoc -> Size
forall (f :: * -> *) vn. [Word8] -> SrcLoc -> ExpBase f vn
StringLit [Word8]
vs SrcLoc
loc
checkExp (IntLit Integer
val NoInfo StructType
NoInfo SrcLoc
loc) = do
  StructType
t <- SrcLoc -> Name -> TermTypeM StructType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
  [PrimType] -> Usage -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
[PrimType] -> Usage -> StructType -> m ()
mustBeOneOf [PrimType]
anyNumberType (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"integer literal") StructType
t
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Integer -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
Integer -> f StructType -> SrcLoc -> ExpBase f vn
IntLit Integer
val (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc
checkExp (FloatLit Double
val NoInfo StructType
NoInfo SrcLoc
loc) = do
  StructType
t <- SrcLoc -> Name -> TermTypeM StructType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
  [PrimType] -> Usage -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
[PrimType] -> Usage -> StructType -> m ()
mustBeOneOf [PrimType]
anyFloatType (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"float literal") StructType
t
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Double -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
Double -> f StructType -> SrcLoc -> ExpBase f vn
FloatLit Double
val (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc
checkExp (TupLit [ExpBase NoInfo VName]
es SrcLoc
loc) =
  [Size] -> SrcLoc -> Size
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Size] -> SrcLoc -> Size)
-> TermTypeM [Size] -> TermTypeM (SrcLoc -> Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase NoInfo VName -> TermTypeM Size)
-> [ExpBase NoInfo VName] -> TermTypeM [Size]
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 ExpBase NoInfo VName -> TermTypeM Size
checkExp [ExpBase NoInfo VName]
es TermTypeM (SrcLoc -> Size) -> TermTypeM SrcLoc -> TermTypeM Size
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 NoInfo VName]
fs SrcLoc
loc) =
  [FieldBase Info VName] -> SrcLoc -> Size
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Size)
-> TermTypeM [FieldBase Info VName] -> TermTypeM (SrcLoc -> Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (Map Name SrcLoc) TermTypeM [FieldBase Info VName]
-> Map Name SrcLoc -> TermTypeM [FieldBase Info VName]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT ((FieldBase NoInfo VName
 -> StateT (Map Name SrcLoc) TermTypeM (FieldBase Info VName))
-> [FieldBase NoInfo VName]
-> StateT (Map Name SrcLoc) 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 NoInfo VName
-> StateT (Map Name SrcLoc) TermTypeM (FieldBase Info VName)
forall {t :: (* -> *) -> * -> *}.
(MonadState (Map Name SrcLoc) (t TermTypeM), MonadTrans t) =>
FieldBase NoInfo VName -> t TermTypeM (FieldBase Info VName)
checkField [FieldBase NoInfo VName]
fs) Map Name SrcLoc
forall a. Monoid a => a
mempty TermTypeM (SrcLoc -> Size) -> TermTypeM SrcLoc -> TermTypeM Size
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 NoInfo VName -> t TermTypeM (FieldBase Info VName)
checkField (RecordFieldExplicit L Name
f ExpBase NoInfo VName
e SrcLoc
rloc) = do
      Name -> SrcLoc -> t TermTypeM ()
forall {a} {b} {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadState (Map a b) (t m), Ord a, MonadTrans t,
 MonadTypeChecker m, Pretty a, Located a, Located b) =>
a -> a -> t m ()
errIfAlreadySet (L Name -> Name
forall a. L a -> a
unLoc L Name
f) SrcLoc
rloc
      (Map Name SrcLoc -> Map Name SrcLoc) -> t TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map Name SrcLoc -> Map Name SrcLoc) -> t TermTypeM ())
-> (Map Name SrcLoc -> Map Name SrcLoc) -> t TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Name -> SrcLoc -> Map Name SrcLoc -> Map Name SrcLoc
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (L Name -> Name
forall a. L a -> a
unLoc L Name
f) SrcLoc
rloc
      L Name -> Size -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit L Name
f (Size -> SrcLoc -> FieldBase Info VName)
-> t TermTypeM Size -> t TermTypeM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TermTypeM Size -> t TermTypeM Size
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e) t TermTypeM (SrcLoc -> FieldBase Info VName)
-> t TermTypeM SrcLoc -> t TermTypeM (FieldBase Info VName)
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
<*> SrcLoc -> t TermTypeM SrcLoc
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
rloc
    checkField (RecordFieldImplicit L VName
name NoInfo StructType
NoInfo SrcLoc
rloc) = do
      Name -> SrcLoc -> t TermTypeM ()
forall {a} {b} {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadState (Map a b) (t m), Ord a, MonadTrans t,
 MonadTypeChecker m, Pretty a, Located a, Located b) =>
a -> a -> t m ()
errIfAlreadySet (VName -> Name
baseName (L VName -> VName
forall a. L a -> a
unLoc L VName
name)) SrcLoc
rloc
      StructType
t <- TermTypeM StructType -> t TermTypeM StructType
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 StructType -> t TermTypeM StructType)
-> TermTypeM StructType -> t TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
rloc (QualName VName -> TermTypeM StructType)
-> QualName VName -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ L VName -> VName
forall a. L a -> a
unLoc L VName
name
      (Map Name SrcLoc -> Map Name SrcLoc) -> t TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map Name SrcLoc -> Map Name SrcLoc) -> t TermTypeM ())
-> (Map Name SrcLoc -> Map Name SrcLoc) -> t TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Name -> SrcLoc -> Map Name SrcLoc -> Map Name SrcLoc
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (VName -> Name
baseName (L VName -> VName
forall a. L a -> a
unLoc L VName
name)) SrcLoc
rloc
      FieldBase Info VName -> t TermTypeM (FieldBase Info VName)
forall a. a -> t TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FieldBase Info VName -> t TermTypeM (FieldBase Info VName))
-> FieldBase Info VName -> t TermTypeM (FieldBase Info VName)
forall a b. (a -> b) -> a -> b
$ L VName -> Info StructType -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L vn -> f StructType -> SrcLoc -> FieldBase f vn
RecordFieldImplicit L VName
name (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
rloc

    errIfAlreadySet :: a -> a -> t m ()
errIfAlreadySet a
f a
rloc = do
      Maybe b
maybe_sloc <- (Map a b -> Maybe b) -> t m (Maybe b)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Map a b -> Maybe b) -> t m (Maybe b))
-> (Map a b -> Maybe b) -> t m (Maybe b)
forall a b. (a -> b) -> a -> b
$ a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
f
      case Maybe b
maybe_sloc of
        Just b
sloc ->
          m () -> t m ()
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 () -> t m ()) -> (Doc () -> m ()) -> Doc () -> t m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> 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 a
rloc Notes
forall a. Monoid a => a
mempty (Doc () -> t m ()) -> Doc () -> t m ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Field"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (a -> Doc ()
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
f)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"previously defined 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 (a -> b -> [Char]
forall a b. (Located a, Located b) => a -> b -> [Char]
locStrRel a
rloc b
sloc)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
        Maybe b
Nothing -> () -> t m ()
forall a. a -> t m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
-- 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) =
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [PrimValue] -> PrimType -> SrcLoc -> Size
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc
checkExp (ArrayLit [ExpBase NoInfo VName]
all_es NoInfo StructType
_ SrcLoc
loc) =
  -- Construct the result type and unify all elements with it.  We
  -- only create a type variable for empty arrays; otherwise we use
  -- the type of the first element.  This significantly cuts down on
  -- the number of type variables generated for pathologically large
  -- multidimensional array literals.
  case [ExpBase NoInfo VName]
all_es of
    [] -> do
      StructType
et <- SrcLoc -> Name -> TermTypeM StructType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
      StructType
t <- SrcLoc -> StructType -> Shape Size -> TermTypeM StructType
arrayOfM SrcLoc
loc StructType
et ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Size
sizeFromInteger Integer
0 SrcLoc
forall a. Monoid a => a
mempty])
      Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [Size] -> Info StructType -> SrcLoc -> Size
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
    ExpBase NoInfo VName
e : [ExpBase NoInfo VName]
es -> do
      Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
      StructType
et <- Size -> TermTypeM StructType
expType Size
e'
      [Size]
es' <- (ExpBase NoInfo VName -> TermTypeM Size)
-> [ExpBase NoInfo VName] -> TermTypeM [Size]
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 (Text -> StructType -> Size -> TermTypeM Size
unifies Text
"type of first array element" StructType
et (Size -> TermTypeM Size)
-> (ExpBase NoInfo VName -> TermTypeM Size)
-> ExpBase NoInfo VName
-> TermTypeM Size
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< ExpBase NoInfo VName -> TermTypeM Size
checkExp) [ExpBase NoInfo VName]
es
      StructType
t <- SrcLoc -> StructType -> Shape Size -> TermTypeM StructType
arrayOfM SrcLoc
loc StructType
et ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Size
sizeFromInteger ([ExpBase NoInfo VName] -> Integer
forall i a. Num i => [a] -> i
genericLength [ExpBase NoInfo VName]
all_es) SrcLoc
forall a. Monoid a => a
mempty])
      Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [Size] -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
ArrayLit (Size
e' Size -> [Size] -> [Size]
forall a. a -> [a] -> [a]
: [Size]
es') (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc
checkExp (AppExp (Range ExpBase NoInfo VName
start Maybe (ExpBase NoInfo VName)
maybe_step Inclusiveness (ExpBase NoInfo VName)
end SrcLoc
loc) NoInfo AppRes
_) = do
  Size
start' <- Text -> [PrimType] -> Size -> TermTypeM Size
require Text
"use in range expression" [PrimType]
anySignedType (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
start
  StructType
start_t <- Size -> TermTypeM StructType
expType Size
start'
  Maybe Size
maybe_step' <- case Maybe (ExpBase NoInfo VName)
maybe_step of
    Maybe (ExpBase NoInfo VName)
Nothing -> Maybe Size -> TermTypeM (Maybe Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Size
forall a. Maybe a
Nothing
    Just ExpBase NoInfo VName
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 (ExpBase NoInfo VName
start, ExpBase NoInfo VName
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 NoInfo StructType
_ SrcLoc
_, Var QualName VName
y_name NoInfo 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
        (ExpBase NoInfo VName, ExpBase NoInfo VName)
_ -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Size -> Maybe Size
forall a. a -> Maybe a
Just (Size -> Maybe Size) -> TermTypeM Size -> TermTypeM (Maybe Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> StructType -> Size -> TermTypeM Size
unifies Text
"use in range expression" StructType
start_t (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
step)

  let unifyRange :: ExpBase NoInfo VName -> TermTypeM Size
unifyRange ExpBase NoInfo VName
e = Text -> StructType -> Size -> TermTypeM Size
unifies Text
"use in range expression" StructType
start_t (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  Inclusiveness Size
end' <- (ExpBase NoInfo VName -> TermTypeM Size)
-> Inclusiveness (ExpBase NoInfo VName)
-> TermTypeM (Inclusiveness Size)
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) -> Inclusiveness a -> f (Inclusiveness b)
traverse ExpBase NoInfo VName -> TermTypeM Size
unifyRange Inclusiveness (ExpBase NoInfo VName)
end

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

  -- Special case some ranges to give them a known size.
  let warnIfBinding :: Bool -> Size -> m (Size, Maybe VName)
warnIfBinding Bool
binds Size
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 (Size -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Size
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."
            VName
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"
            (Size, Maybe VName) -> m (Size, Maybe VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> SrcLoc -> Size
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
d) SrcLoc
forall a. Monoid a => a
mempty, VName -> Maybe VName
forall a. a -> Maybe a
Just VName
d)
          else (Size, Maybe VName) -> m (Size, Maybe VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size
size, Maybe VName
forall a. Maybe a
Nothing)
  (Size
dim, Maybe VName
retext) <-
    case (Size -> Maybe Int64
isInt64 Size
start', Size -> Maybe Int64
isInt64 (Size -> Maybe Int64) -> Maybe Size -> Maybe (Maybe Int64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Size
maybe_step', Inclusiveness Size
end') of
      (Just Int64
0, Just (Just Int64
1), UpToExclusive Size
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Size -> TermTypeM (Size, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Size -> m (Size, Maybe VName)
warnIfBinding (Size -> Bool
hasBinding Size
end'') Size
end''
      (Just Int64
0, Maybe (Maybe Int64)
Nothing, UpToExclusive Size
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Size -> TermTypeM (Size, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Size -> m (Size, Maybe VName)
warnIfBinding (Size -> Bool
hasBinding Size
end'') Size
end''
      (Maybe Int64
_, Maybe (Maybe Int64)
Nothing, UpToExclusive Size
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Size -> TermTypeM (Size, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Size -> m (Size, Maybe VName)
warnIfBinding (Size -> Bool
hasBinding Size
end'' Bool -> Bool -> Bool
|| Size -> Bool
hasBinding Size
start') (Size -> TermTypeM (Size, Maybe VName))
-> Size -> TermTypeM (Size, Maybe VName)
forall a b. (a -> b) -> a -> b
$ Size -> Size -> Size
sizeMinus Size
end'' Size
start'
      (Maybe Int64
_, Maybe (Maybe Int64)
Nothing, ToInclusive Size
end'')
        -- No stride means we assume a stride of one.
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Size -> TermTypeM (Size, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Size -> m (Size, Maybe VName)
warnIfBinding (Size -> Bool
hasBinding Size
end'' Bool -> Bool -> Bool
|| Size -> Bool
hasBinding Size
start') (Size -> TermTypeM (Size, Maybe VName))
-> Size -> TermTypeM (Size, Maybe VName)
forall a b. (a -> b) -> a -> b
$ Size -> Size -> Size
sizeMinusInc Size
end'' Size
start'
      (Just Int64
1, Just (Just Int64
2), ToInclusive Size
end'')
        | Scalar (Prim (Signed IntType
Int64)) <- StructType
end_t ->
            Bool -> Size -> TermTypeM (Size, Maybe VName)
forall {m :: * -> *}.
(MonadTypeChecker m, MonadUnify m) =>
Bool -> Size -> m (Size, Maybe VName)
warnIfBinding (Size -> Bool
hasBinding Size
end'') Size
end''
      (Maybe Int64, Maybe (Maybe Int64), Inclusiveness Size)
_ -> do
        VName
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"
        (Size, Maybe VName) -> TermTypeM (Size, Maybe VName)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> SrcLoc -> Size
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
d) SrcLoc
forall a. Monoid a => a
mempty, VName -> Maybe VName
forall a. a -> Maybe a
Just VName
d)

  StructType
t <- SrcLoc -> StructType -> Shape Size -> TermTypeM StructType
arrayOfM SrcLoc
loc StructType
start_t ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [Size
dim])
  let res :: AppRes
res = StructType -> [VName] -> AppRes
AppRes StructType
t (Maybe VName -> [VName]
forall a. Maybe a -> [a]
maybeToList Maybe VName
retext)

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Size
-> Maybe Size
-> Inclusiveness Size
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Size
start' Maybe Size
maybe_step' Inclusiveness Size
end' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
  where
    i64 :: TypeBase dim u
i64 = ScalarTypeBase dim u -> TypeBase dim u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim u -> TypeBase dim u)
-> ScalarTypeBase dim u -> TypeBase dim u
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase dim u)
-> PrimType -> ScalarTypeBase dim u
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    mkBinOp :: Name -> StructType -> Size -> Size -> Size
mkBinOp Name
op StructType
t Size
x Size
y =
      AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        ( (QualName VName, SrcLoc)
-> Info StructType
-> (Size, Info (Maybe VName))
-> (Size, 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
sizeBinOpInfo
            (Size
x, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
            (Size
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 :: Size -> Size -> Size
mkSub = Name -> StructType -> Size -> Size -> Size
mkBinOp Name
"-" StructType
forall {dim} {u}. TypeBase dim u
i64
    mkAdd :: Size -> Size -> Size
mkAdd = Name -> StructType -> Size -> Size -> Size
mkBinOp Name
"+" StructType
forall {dim} {u}. TypeBase dim u
i64
    sizeMinus :: Size -> Size -> Size
sizeMinus Size
j Size
i = Size
j Size -> Size -> Size
`mkSub` Size
i
    sizeMinusInc :: Size -> Size -> Size
sizeMinusInc Size
j Size
i = (Size
j Size -> Size -> Size
`mkSub` Size
i) Size -> Size -> Size
`mkAdd` Integer -> SrcLoc -> Size
sizeFromInteger Integer
1 SrcLoc
forall a. Monoid a => a
mempty
    sizeBinOpInfo :: Info StructType
sizeBinOpInfo = StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ [ParamType] -> ResRetType -> StructType
foldFunType [ParamType
forall {dim} {u}. TypeBase dim u
i64, ParamType
forall {dim} {u}. TypeBase dim u
i64] (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Size Uniqueness
forall {dim} {u}. TypeBase dim u
i64
checkExp (Ascript ExpBase NoInfo VName
e TypeExp (ExpBase NoInfo VName) VName
te SrcLoc
loc) = do
  (TypeExp Size VName
te', Size
e') <- SrcLoc
-> TypeExp (ExpBase NoInfo VName) VName
-> ExpBase NoInfo VName
-> TermTypeM (TypeExp Size VName, Size)
checkAscript SrcLoc
loc TypeExp (ExpBase NoInfo VName) VName
te ExpBase NoInfo VName
e
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> TypeExp Size VName -> SrcLoc -> Size
forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp (ExpBase f vn) vn -> SrcLoc -> ExpBase f vn
Ascript Size
e' TypeExp Size VName
te' SrcLoc
loc
checkExp (Coerce ExpBase NoInfo VName
e TypeExp (ExpBase NoInfo VName) VName
te NoInfo StructType
NoInfo SrcLoc
loc) = do
  (TypeExp Size VName
te', StructType
te_t, Size
e') <- SrcLoc
-> TypeExp (ExpBase NoInfo VName) VName
-> ExpBase NoInfo VName
-> TermTypeM (TypeExp Size VName, StructType, Size)
checkCoerce SrcLoc
loc TypeExp (ExpBase NoInfo VName) VName
te ExpBase NoInfo VName
e
  StructType
t <- Size -> TermTypeM StructType
expTypeFully Size
e'
  StructType
t' <- ([VName] -> Size -> Size -> TermTypeM Size)
-> StructType -> StructType -> TermTypeM StructType
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims ((Size -> TermTypeM Size) -> Size -> Size -> TermTypeM Size
forall a b. a -> b -> a
const ((Size -> TermTypeM Size) -> Size -> Size -> TermTypeM Size)
-> ([VName] -> Size -> TermTypeM Size)
-> [VName]
-> Size
-> Size
-> TermTypeM Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Size -> TermTypeM Size) -> [VName] -> Size -> TermTypeM Size
forall a b. a -> b -> a
const Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) StructType
t StructType
te_t
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> TypeExp Size VName -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp (ExpBase f vn) vn
-> f StructType
-> SrcLoc
-> ExpBase f vn
Coerce Size
e' TypeExp Size VName
te' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t') SrcLoc
loc
checkExp (AppExp (BinOp (QualName VName
op, SrcLoc
oploc) NoInfo StructType
NoInfo (ExpBase NoInfo VName
e1, NoInfo (Maybe VName)
_) (ExpBase NoInfo VName
e2, NoInfo (Maybe VName)
_) SrcLoc
loc) NoInfo AppRes
NoInfo) = do
  StructType
ftype <- SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
oploc QualName VName
op
  Size
e1' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e1
  Size
e2' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e2

  -- Note that the application to the first operand cannot fix any
  -- existential sizes, because it must by necessity be a function.
  (StructType
_, StructType
rt, Maybe VName
p1_ext, [VName]
_) <- SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
op, Int
0) StructType
ftype Size
e1'
  (StructType
_, StructType
rt', Maybe VName
p2_ext, [VName]
retext) <- SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
op, Int
1) StructType
rt Size
e2'

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
    AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
      ( (QualName VName, SrcLoc)
-> Info StructType
-> (Size, Info (Maybe VName))
-> (Size, 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
          (QualName VName
op, SrcLoc
oploc)
          (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ftype)
          (Size
e1', Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
p1_ext)
          (Size
e2', Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
p2_ext)
          SrcLoc
loc
      )
      (AppRes -> Info AppRes
forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes StructType
rt' [VName]
retext))
checkExp (Project Name
k ExpBase NoInfo VName
e NoInfo StructType
NoInfo SrcLoc
loc) = do
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  StructType
t <- Size -> TermTypeM StructType
expType Size
e'
  StructType
kt <- Usage -> Name -> StructType -> TermTypeM StructType
forall (m :: * -> *).
MonadUnify m =>
Usage -> Name -> StructType -> m StructType
mustHaveField (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc (Text -> Usage) -> Text -> Usage
forall a b. (a -> b) -> a -> b
$ Doc Any -> Text
forall a. Doc a -> Text
docText (Doc Any -> Text) -> Doc Any -> Text
forall a b. (a -> b) -> a -> b
$ Doc Any
"projection of field " Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann
dquotes (Name -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. Name -> Doc ann
pretty Name
k)) Name
k StructType
t
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Name -> Size -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
k Size
e' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
kt) SrcLoc
loc
checkExp (AppExp (If ExpBase NoInfo VName
e1 ExpBase NoInfo VName
e2 ExpBase NoInfo VName
e3 SrcLoc
loc) NoInfo AppRes
_) = do
  Size
e1' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e1
  Size
e2' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e2
  Size
e3' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e3

  let bool :: TypeBase dim u
bool = ScalarTypeBase dim u -> TypeBase dim u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim u -> TypeBase dim u)
-> ScalarTypeBase dim u -> TypeBase dim u
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
Bool
  StructType
e1_t <- Size -> TermTypeM StructType
expType Size
e1'
  Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure ([StructType] -> StructType -> Checking
CheckingRequired [StructType
forall {dim} {u}. TypeBase dim u
bool] StructType
e1_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 (Size -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Size
e1' Text
"use as 'if' condition") StructType
forall {dim} {u}. TypeBase dim u
bool StructType
e1_t

  (StructType
brancht, [VName]
retext) <- SrcLoc -> Size -> Size -> TermTypeM (StructType, [VName])
unifyBranches SrcLoc
loc Size
e2' Size
e3'

  Usage -> Text -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> Text -> StructType -> m ()
zeroOrderType
    (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"returning value of this type from 'if' expression")
    Text
"type returned from branch"
    StructType
brancht

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Size -> Size -> Size -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Size
e1' Size
e2' Size
e3' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
brancht [VName]
retext)
checkExp (Parens ExpBase NoInfo VName
e SrcLoc
loc) =
  Size -> SrcLoc -> Size
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Size -> SrcLoc -> Size)
-> TermTypeM Size -> TermTypeM (SrcLoc -> Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e TermTypeM (SrcLoc -> Size) -> TermTypeM SrcLoc -> TermTypeM Size
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) ExpBase NoInfo VName
e SrcLoc
loc) = do
  Mod
mod <- QualName VName -> TermTypeM Mod
lookupMod QualName VName
modname
  case Mod
mod of
    ModEnv Env
env -> (TermEnv -> TermEnv) -> TermTypeM Size -> TermTypeM Size
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 Size -> TermTypeM Size)
-> TermTypeM Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ do
      Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
      Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ (QualName VName, SrcLoc) -> Size -> SrcLoc -> Size
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName
modname, SrcLoc
modnameloc) Size
e' SrcLoc
loc
    ModFun {} ->
      SrcLoc -> Notes -> Doc () -> TermTypeM Size
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 Size)
-> (Doc () -> Doc ()) -> Doc () -> TermTypeM Size
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 Size) -> Doc () -> TermTypeM Size
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 (Var QualName VName
qn NoInfo StructType
NoInfo SrcLoc
loc) = do
  StructType
t <- SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
qn
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc
checkExp (Negate ExpBase NoInfo VName
arg SrcLoc
loc) = do
  Size
arg' <- Text -> [PrimType] -> Size -> TermTypeM Size
require Text
"numeric negation" [PrimType]
anyNumberType (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
arg
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> SrcLoc -> Size
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate Size
arg' SrcLoc
loc
checkExp (Not ExpBase NoInfo VName
arg SrcLoc
loc) = do
  Size
arg' <- Text -> [PrimType] -> Size -> TermTypeM Size
require Text
"logical negation" (PrimType
Bool PrimType -> [PrimType] -> [PrimType]
forall a. a -> [a] -> [a]
: [PrimType]
anyIntType) (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
arg
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> SrcLoc -> Size
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not Size
arg' SrcLoc
loc
checkExp (AppExp (Apply ExpBase NoInfo VName
fe NonEmpty (NoInfo (Maybe VName), ExpBase NoInfo VName)
args SrcLoc
loc) NoInfo AppRes
NoInfo) = do
  Size
fe' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
fe
  NonEmpty Size
args' <- ((NoInfo (Maybe VName), ExpBase NoInfo VName) -> TermTypeM Size)
-> NonEmpty (NoInfo (Maybe VName), ExpBase NoInfo VName)
-> TermTypeM (NonEmpty Size)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM (ExpBase NoInfo VName -> TermTypeM Size
checkExp (ExpBase NoInfo VName -> TermTypeM Size)
-> ((NoInfo (Maybe VName), ExpBase NoInfo VName)
    -> ExpBase NoInfo VName)
-> (NoInfo (Maybe VName), ExpBase NoInfo VName)
-> TermTypeM Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NoInfo (Maybe VName), ExpBase NoInfo VName)
-> ExpBase NoInfo VName
forall a b. (a, b) -> b
snd) NonEmpty (NoInfo (Maybe VName), ExpBase NoInfo VName)
args
  StructType
t <- Size -> TermTypeM StructType
expType Size
fe'
  let fname :: Maybe (QualName VName)
fname =
        case Size
fe' of
          Var QualName VName
v Info StructType
_ SrcLoc
_ -> QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
v
          Size
_ -> Maybe (QualName VName)
forall a. Maybe a
Nothing
  ((Int
_, [VName]
exts, StructType
rt), NonEmpty (Info (Maybe VName), Size)
args'') <- ((Int, [VName], StructType)
 -> Size
 -> TermTypeM
      ((Int, [VName], StructType), (Info (Maybe VName), Size)))
-> (Int, [VName], StructType)
-> NonEmpty Size
-> TermTypeM
     ((Int, [VName], StructType), NonEmpty (Info (Maybe VName), Size))
forall (m :: * -> *) (t :: * -> *) acc x y.
(Monad m, Traversable t) =>
(acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
mapAccumLM (Maybe (QualName VName)
-> (Int, [VName], StructType)
-> Size
-> TermTypeM
     ((Int, [VName], StructType), (Info (Maybe VName), Size))
onArg Maybe (QualName VName)
fname) (Int
0, [], StructType
t) NonEmpty Size
args'

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Size
-> NonEmpty (Info (Maybe VName), Size)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Size
fe' NonEmpty (Info (Maybe VName), Size)
args'' SrcLoc
loc) (Info AppRes -> Size) -> Info AppRes -> Size
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
rt [VName]
exts
  where
    onArg :: Maybe (QualName VName)
-> (Int, [VName], StructType)
-> Size
-> TermTypeM
     ((Int, [VName], StructType), (Info (Maybe VName), Size))
onArg Maybe (QualName VName)
fname (Int
i, [VName]
all_exts, StructType
t) Size
arg' = do
      (StructType
_, StructType
rt, Maybe VName
argext, [VName]
exts) <- SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (Maybe (QualName VName)
fname, Int
i) StructType
t Size
arg'
      ((Int, [VName], StructType), (Info (Maybe VName), Size))
-> TermTypeM
     ((Int, [VName], StructType), (Info (Maybe VName), Size))
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, [VName]
all_exts [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
exts, StructType
rt),
          (Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
argext, Size
arg')
        )
checkExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase NoInfo VName StructType
pat ExpBase NoInfo VName
e ExpBase NoInfo VName
body SrcLoc
loc) NoInfo AppRes
_) = do
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e

  -- Not technically an ascription, but we want the pattern to have
  -- exactly the type of 'e'.
  StructType
t <- Size -> TermTypeM StructType
expType Size
e'
  [SizeBinder VName] -> TermTypeM Size -> TermTypeM Size
forall a. [SizeBinder VName] -> TermTypeM a -> TermTypeM a
bindingSizes [SizeBinder VName]
sizes (TermTypeM Size -> TermTypeM Size)
-> ((Pat ParamType -> TermTypeM Size) -> TermTypeM Size)
-> (Pat ParamType -> TermTypeM Size)
-> TermTypeM Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeM Size -> TermTypeM Size
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM Size -> TermTypeM Size)
-> ((Pat ParamType -> TermTypeM Size) -> TermTypeM Size)
-> (Pat ParamType -> TermTypeM Size)
-> TermTypeM Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SizeBinder VName]
-> PatBase NoInfo VName StructType
-> StructType
-> (Pat ParamType -> TermTypeM Size)
-> TermTypeM Size
forall u a.
[SizeBinder VName]
-> PatBase NoInfo VName (TypeBase Size u)
-> StructType
-> (Pat ParamType -> TermTypeM a)
-> TermTypeM a
bindingPat [SizeBinder VName]
sizes PatBase NoInfo VName StructType
pat StructType
t ((Pat ParamType -> TermTypeM Size) -> TermTypeM Size)
-> (Pat ParamType -> TermTypeM Size) -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ \Pat ParamType
pat' -> do
    Size
body' <- TermTypeM Size -> TermTypeM Size
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM Size -> TermTypeM Size)
-> TermTypeM Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
body
    StructType
body_t <- Size -> TermTypeM StructType
expTypeFully Size
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.
    StructType
t' <- StructType -> TermTypeM StructType
forall (m :: * -> *). MonadUnify m => StructType -> m StructType
normType StructType
t -- Might be overloaded integer until now.
    (StructType
body_t', [VName]
retext) <-
      case (StructType
t', Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames Pat ParamType
pat') of
        (Scalar (Prim (Signed IntType
Int64)), [VName
v])
          | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Size -> Bool
hasBinding Size
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 (Size -> Subst t
forall t. Size -> Subst t
ExpSubst Size
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 as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [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

    Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
      AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        ([SizeBinder VName]
-> PatBase Info VName StructType
-> Size
-> Size
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes ((ParamType -> StructType)
-> Pat ParamType -> PatBase Info VName StructType
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct Pat ParamType
pat') Size
e' Size
body' SrcLoc
loc)
        (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
body_t' [VName]
retext)
checkExp (AppExp (LetFun VName
name ([TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl, NoInfo ResRetType
NoInfo, ExpBase NoInfo VName
e) ExpBase NoInfo VName
body SrcLoc
loc) NoInfo AppRes
_) = do
  ([TypeParamBase VName]
tparams', [Pat ParamType]
params', Maybe (TypeExp Size VName)
maybe_retdecl', ResRetType
rettype, Size
e') <-
    (VName, Maybe (TypeExp (ExpBase NoInfo VName) VName),
 [TypeParamBase VName], [PatBase NoInfo VName ParamType],
 ExpBase NoInfo VName, SrcLoc)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
checkBinding (VName
name, Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl, [TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, ExpBase NoInfo VName
e, SrcLoc
loc)

  let entry :: ValBinding
entry = [TypeParamBase VName] -> StructType -> ValBinding
BoundV [TypeParamBase VName]
tparams' (StructType -> ValBinding) -> StructType -> ValBinding
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> ResRetType -> StructType
funType [Pat ParamType]
params' ResRetType
rettype
      bindF :: TermScope -> TermScope
bindF TermScope
scope =
        TermScope
scope
          { scopeVtable = M.insert name entry $ scopeVtable scope
          }
  Size
body' <- (TermScope -> TermScope) -> TermTypeM Size -> TermTypeM Size
forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope TermScope -> TermScope
bindF (TermTypeM Size -> TermTypeM Size)
-> TermTypeM Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
body

  (StructType
body_t, [VName]
ext) <- SrcLoc -> [VName] -> StructType -> TermTypeM (StructType, [VName])
forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType SrcLoc
loc [VName
name] (StructType -> TermTypeM (StructType, [VName]))
-> TermTypeM StructType -> TermTypeM (StructType, [VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Size -> TermTypeM StructType
expTypeFully Size
body'

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
    AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
      ( VName
-> ([TypeParamBase VName], [Pat ParamType],
    Maybe (TypeExp Size VName), Info ResRetType, Size)
-> Size
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn ParamType],
    Maybe (TypeExp (ExpBase f vn) vn), f ResRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun
          VName
name
          ([TypeParamBase VName]
tparams', [Pat ParamType]
params', Maybe (TypeExp Size VName)
maybe_retdecl', ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
rettype, Size
e')
          Size
body'
          SrcLoc
loc
      )
      (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
body_t [VName]
ext)
checkExp (AppExp (LetWith IdentBase NoInfo VName StructType
dest IdentBase NoInfo VName StructType
src SliceBase NoInfo VName
slice ExpBase NoInfo VName
ve ExpBase NoInfo VName
body SrcLoc
loc) NoInfo AppRes
_) = do
  Ident StructType
src' <- IdentBase NoInfo VName StructType -> TermTypeM (Ident StructType)
checkIdent IdentBase NoInfo VName StructType
src
  [DimIndex]
slice' <- SliceBase NoInfo VName -> TermTypeM [DimIndex]
checkSlice SliceBase NoInfo VName
slice
  (StructType
t, StructType
_) <- Usage -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType (IdentBase NoInfo VName StructType -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage IdentBase NoInfo VName StructType
src Text
"type of source array") Name
"src" (Int -> TermTypeM (StructType, StructType))
-> Int -> TermTypeM (StructType, StructType)
forall a b. (a -> b) -> a -> b
$ [DimIndex] -> Int
sliceDims [DimIndex]
slice'
  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
"type of target array") StructType
t (StructType -> TermTypeM ()) -> StructType -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> StructType) -> Info StructType -> StructType
forall a b. (a -> b) -> a -> b
$ Ident StructType -> Info StructType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType Ident StructType
src'

  (StructType
elemt, [VName]
_) <- Maybe (SrcLoc, Rigidity)
-> [DimIndex] -> StructType -> TermTypeM (StructType, [VName])
forall as.
Maybe (SrcLoc, Rigidity)
-> [DimIndex]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
sliceShape ((SrcLoc, Rigidity) -> Maybe (SrcLoc, Rigidity)
forall a. a -> Maybe a
Just (SrcLoc
loc, Rigidity
Nonrigid)) [DimIndex]
slice' (StructType -> TermTypeM (StructType, [VName]))
-> TermTypeM StructType -> TermTypeM (StructType, [VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< StructType -> TermTypeM StructType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully StructType
t

  Size
ve' <- Text -> StructType -> Size -> TermTypeM Size
unifies Text
"type of target array" StructType
elemt (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
ve

  IdentBase NoInfo VName StructType
-> StructType
-> (Ident StructType -> TermTypeM Size)
-> TermTypeM Size
forall a.
IdentBase NoInfo VName StructType
-> StructType -> (Ident StructType -> TermTypeM a) -> TermTypeM a
bindingIdent IdentBase NoInfo VName StructType
dest (Info StructType -> StructType
forall a. Info a -> a
unInfo (Ident StructType -> Info StructType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType Ident StructType
src')) ((Ident StructType -> TermTypeM Size) -> TermTypeM Size)
-> (Ident StructType -> TermTypeM Size) -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ \Ident StructType
dest' -> do
    Size
body' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
body
    (StructType
body_t, [VName]
ext) <- SrcLoc -> [VName] -> StructType -> TermTypeM (StructType, [VName])
forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType SrcLoc
loc [Ident StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName Ident StructType
dest'] (StructType -> TermTypeM (StructType, [VName]))
-> TermTypeM StructType -> TermTypeM (StructType, [VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Size -> TermTypeM StructType
expTypeFully Size
body'
    Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Ident StructType
-> Ident StructType
-> [DimIndex]
-> Size
-> Size
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn StructType
-> IdentBase f vn StructType
-> SliceBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetWith Ident StructType
dest' Ident StructType
src' [DimIndex]
slice' Size
ve' Size
body' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
body_t [VName]
ext)
checkExp (Update ExpBase NoInfo VName
src SliceBase NoInfo VName
slice ExpBase NoInfo VName
ve SrcLoc
loc) = do
  [DimIndex]
slice' <- SliceBase NoInfo VName -> TermTypeM [DimIndex]
checkSlice SliceBase NoInfo VName
slice
  (StructType
t, StructType
_) <- Usage -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType (ExpBase NoInfo VName -> Usage
forall a. Located a => a -> Usage
mkUsage' ExpBase NoInfo VName
src) Name
"src" (Int -> TermTypeM (StructType, StructType))
-> Int -> TermTypeM (StructType, StructType)
forall a b. (a -> b) -> a -> b
$ [DimIndex] -> Int
sliceDims [DimIndex]
slice'
  (StructType
elemt, [VName]
_) <- Maybe (SrcLoc, Rigidity)
-> [DimIndex] -> StructType -> TermTypeM (StructType, [VName])
forall as.
Maybe (SrcLoc, Rigidity)
-> [DimIndex]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
sliceShape ((SrcLoc, Rigidity) -> Maybe (SrcLoc, Rigidity)
forall a. a -> Maybe a
Just (SrcLoc
loc, Rigidity
Nonrigid)) [DimIndex]
slice' (StructType -> TermTypeM (StructType, [VName]))
-> TermTypeM StructType -> TermTypeM (StructType, [VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< StructType -> TermTypeM StructType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully StructType
t
  Size
ve' <- Text -> StructType -> Size -> TermTypeM Size
unifies Text
"type of target array" StructType
elemt (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
ve
  Size
src' <- Text -> StructType -> Size -> TermTypeM Size
unifies Text
"type of target array" StructType
t (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
src
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> [DimIndex] -> Size -> SrcLoc -> Size
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update Size
src' [DimIndex]
slice' Size
ve' SrcLoc
loc

-- 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 (RecordUpdate ExpBase NoInfo VName
src [Name]
fields ExpBase NoInfo VName
ve NoInfo StructType
NoInfo SrcLoc
loc) = do
  Size
src' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
src
  Size
ve' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
ve
  StructType
a <- Size -> TermTypeM StructType
expTypeFully Size
src'
  (StructType -> Name -> TermTypeM StructType)
-> StructType -> [Name] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ ((Name -> StructType -> TermTypeM StructType)
-> StructType -> Name -> TermTypeM StructType
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Name -> StructType -> TermTypeM StructType)
 -> StructType -> Name -> TermTypeM StructType)
-> (Name -> StructType -> TermTypeM StructType)
-> StructType
-> Name
-> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ Usage -> Name -> StructType -> TermTypeM StructType
forall (m :: * -> *).
MonadUnify m =>
Usage -> Name -> StructType -> m StructType
mustHaveField Usage
usage) StructType
a [Name]
fields
  StructType
ve_t <- Size -> TermTypeM StructType
expType Size
ve'
  StructType
updated_t <- [Name] -> StructType -> StructType -> TermTypeM StructType
updateField [Name]
fields StructType
ve_t (StructType -> TermTypeM StructType)
-> TermTypeM StructType -> TermTypeM StructType
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Size -> TermTypeM StructType
expTypeFully Size
src'
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> [Name] -> Size -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
RecordUpdate Size
src' [Name]
fields Size
ve' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
updated_t) SrcLoc
loc
  where
    usage :: Usage
usage = SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"record update"
    updateField :: [Name] -> StructType -> StructType -> TermTypeM StructType
updateField [] StructType
ve_t StructType
src_t = do
      (StructType
src_t', Map VName Size
_) <- Usage
-> Rigidity
-> Name
-> StructType
-> TermTypeM (StructType, Map VName Size)
forall als.
Usage
-> Rigidity
-> Name
-> TypeBase Size als
-> TermTypeM (TypeBase Size als, Map VName Size)
allDimsFreshInType Usage
usage Rigidity
Nonrigid Name
"any" StructType
src_t
      Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure ([Name] -> StructType -> StructType -> Checking
CheckingRecordUpdate [Name]
fields StructType
src_t' StructType
ve_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 StructType
src_t' StructType
ve_t
      StructType -> TermTypeM StructType
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StructType
ve_t
    updateField (Name
f : [Name]
fs) StructType
ve_t (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
          StructType
f_t' <- [Name] -> StructType -> StructType -> TermTypeM StructType
updateField [Name]
fs StructType
ve_t StructType
f_t
          StructType -> TermTypeM StructType
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType -> TermTypeM StructType)
-> StructType -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size NoUniqueness -> StructType)
-> ScalarTypeBase Size NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Size NoUniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name StructType -> ScalarTypeBase Size NoUniqueness)
-> Map Name StructType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ Name -> StructType -> Map Name StructType -> Map Name StructType
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
f StructType
f_t' Map Name StructType
m
    updateField [Name]
_ StructType
_ StructType
_ =
      SrcLoc -> 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 SrcLoc
loc 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
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ExpBase NoInfo VName -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ExpBase NoInfo VName -> Doc ann
pretty ExpBase NoInfo VName
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."

--
checkExp (AppExp (Index ExpBase NoInfo VName
e SliceBase NoInfo VName
slice SrcLoc
loc) NoInfo AppRes
_) = do
  [DimIndex]
slice' <- SliceBase NoInfo VName -> TermTypeM [DimIndex]
checkSlice SliceBase NoInfo VName
slice
  (StructType
t, StructType
_) <- Usage -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType (SrcLoc -> Usage
forall a. Located a => a -> Usage
mkUsage' SrcLoc
loc) Name
"e" (Int -> TermTypeM (StructType, StructType))
-> Int -> TermTypeM (StructType, StructType)
forall a b. (a -> b) -> a -> b
$ [DimIndex] -> Int
sliceDims [DimIndex]
slice'
  Size
e' <- Text -> StructType -> Size -> TermTypeM Size
unifies Text
"being indexed at" StructType
t (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value.
  (StructType
t', [VName]
retext) <-
    Maybe (SrcLoc, Rigidity)
-> [DimIndex] -> StructType -> TermTypeM (StructType, [VName])
forall as.
Maybe (SrcLoc, Rigidity)
-> [DimIndex]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
sliceShape ((SrcLoc, Rigidity) -> Maybe (SrcLoc, Rigidity)
forall a. a -> Maybe a
Just (SrcLoc
loc, RigidSource -> Rigidity
Rigid (Maybe Size -> Text -> RigidSource
RigidSlice Maybe Size
forall a. Maybe a
Nothing Text
""))) [DimIndex]
slice'
      (StructType -> TermTypeM (StructType, [VName]))
-> TermTypeM StructType -> TermTypeM (StructType, [VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Size -> TermTypeM StructType
expTypeFully Size
e'

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Size -> [DimIndex] -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index Size
e' [DimIndex]
slice' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
t' [VName]
retext)
checkExp (Assert ExpBase NoInfo VName
e1 ExpBase NoInfo VName
e2 NoInfo Text
NoInfo SrcLoc
loc) = do
  Size
e1' <- Text -> [PrimType] -> Size -> TermTypeM Size
require Text
"being asserted" [PrimType
Bool] (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e1
  Size
e2' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e2
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Size -> Size -> Info Text -> SrcLoc -> Size
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert Size
e1' Size
e2' (Text -> Info Text
forall a. a -> Info a
Info (ExpBase NoInfo VName -> Text
forall a. Pretty a => a -> Text
prettyText ExpBase NoInfo VName
e1)) SrcLoc
loc
checkExp (Lambda [PatBase NoInfo VName ParamType]
params ExpBase NoInfo VName
body Maybe (TypeExp (ExpBase NoInfo VName) VName)
rettype_te NoInfo ResRetType
NoInfo SrcLoc
loc) = do
  ([Pat ParamType]
params', Size
body', Maybe (TypeExp Size VName)
rettype', RetType [VName]
dims TypeBase Size Uniqueness
ty) <-
    TermTypeM
  ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
-> TermTypeM
     ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM
   ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
 -> TermTypeM
      ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
-> (([Pat ParamType]
     -> TermTypeM
          ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
    -> TermTypeM
         ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
-> ([Pat ParamType]
    -> TermTypeM
         ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
-> TermTypeM
     ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeParamBase VName]
-> [PatBase NoInfo VName ParamType]
-> ([Pat ParamType]
    -> TermTypeM
         ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
-> TermTypeM
     ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
forall a.
[TypeParamBase VName]
-> [PatBase NoInfo VName ParamType]
-> ([Pat ParamType] -> TermTypeM a)
-> TermTypeM a
bindingParams [] [PatBase NoInfo VName ParamType]
params (([Pat ParamType]
  -> TermTypeM
       ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
 -> TermTypeM
      ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
-> ([Pat ParamType]
    -> TermTypeM
         ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType))
-> TermTypeM
     ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
      Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
rettype_checked <- (TypeExp (ExpBase NoInfo VName) VName
 -> TermTypeM
      (TypeExp Size VName, TypeBase Size Uniqueness, [VName]))
-> Maybe (TypeExp (ExpBase NoInfo VName) VName)
-> TermTypeM
     (Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [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 (ExpBase NoInfo VName) VName
-> TermTypeM
     (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
checkTypeExpNonrigid Maybe (TypeExp (ExpBase NoInfo VName) VName)
rettype_te
      let declared_rettype :: Maybe (TypeBase Size Uniqueness)
declared_rettype =
            case Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
rettype_checked of
              Just (TypeExp Size VName
_, TypeBase Size Uniqueness
st, [VName]
_) -> TypeBase Size Uniqueness -> Maybe (TypeBase Size Uniqueness)
forall a. a -> Maybe a
Just TypeBase Size Uniqueness
st
              Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
Nothing -> Maybe (TypeBase Size Uniqueness)
forall a. Maybe a
Nothing
      Size
body' <- [Pat ParamType]
-> ExpBase NoInfo VName
-> Maybe (TypeBase Size Uniqueness)
-> SrcLoc
-> TermTypeM Size
checkFunBody [Pat ParamType]
params' ExpBase NoInfo VName
body Maybe (TypeBase Size Uniqueness)
declared_rettype SrcLoc
loc
      StructType
body_t <- Size -> TermTypeM StructType
expTypeFully Size
body'

      [Pat ParamType]
params'' <- (Pat ParamType -> TermTypeM (Pat ParamType))
-> [Pat ParamType] -> TermTypeM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> TermTypeM (Pat ParamType)
forall e. ASTMappable e => e -> TermTypeM e
updateTypes [Pat ParamType]
params'

      (Maybe (TypeExp Size VName)
rettype', ResRetType
rettype_st) <-
        case Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
rettype_checked of
          Just (TypeExp Size VName
te, TypeBase Size Uniqueness
st, [VName]
ext) ->
            (Maybe (TypeExp Size VName), ResRetType)
-> TermTypeM (Maybe (TypeExp Size VName), ResRetType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp Size VName -> Maybe (TypeExp Size VName)
forall a. a -> Maybe a
Just TypeExp Size VName
te, [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext TypeBase Size Uniqueness
st)
          Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
Nothing -> do
            ResRetType
ret <- [Pat ParamType] -> TypeBase Size Uniqueness -> TermTypeM ResRetType
forall {m :: * -> *}.
MonadUnify m =>
[Pat ParamType] -> TypeBase Size Uniqueness -> m ResRetType
inferReturnSizes [Pat ParamType]
params'' (TypeBase Size Uniqueness -> TermTypeM ResRetType)
-> TypeBase Size Uniqueness -> TermTypeM ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> TypeBase Size Uniqueness
forall u. Uniqueness -> TypeBase Size u -> TypeBase Size Uniqueness
toRes Uniqueness
Nonunique StructType
body_t
            (Maybe (TypeExp Size VName), ResRetType)
-> TermTypeM (Maybe (TypeExp Size VName), ResRetType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (TypeExp Size VName)
forall a. Maybe a
Nothing, ResRetType
ret)

      ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
-> TermTypeM
     ([Pat ParamType], Size, Maybe (TypeExp Size VName), ResRetType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Pat ParamType]
params'', Size
body', Maybe (TypeExp Size VName)
rettype', ResRetType
rettype_st)

  Maybe VName -> [Pat ParamType] -> TermTypeM ()
verifyFunctionParams Maybe VName
forall a. Maybe a
Nothing [Pat ParamType]
params'

  (TypeBase Size Uniqueness
ty', [VName]
dims') <- SrcLoc
-> [VName]
-> TypeBase Size Uniqueness
-> TermTypeM (TypeBase Size Uniqueness, [VName])
forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType SrcLoc
loc [VName]
dims TypeBase Size Uniqueness
ty

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [Pat ParamType]
-> Size
-> Maybe (TypeExp Size VName)
-> Info ResRetType
-> SrcLoc
-> Size
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
params' Size
body' Maybe (TypeExp Size VName)
rettype' (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims' TypeBase Size Uniqueness
ty')) SrcLoc
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] -> TypeBase Size Uniqueness -> m ResRetType
inferReturnSizes [Pat ParamType]
params' TypeBase Size Uniqueness
ret = do
      Int
cur_lvl <- m Int
forall (m :: * -> *). MonadUnify m => m Int
curLevel
      let 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 :: [VName]
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)
patternParam) [Pat ParamType]
params'
          pos_sizes :: Set VName
pos_sizes =
            StructType -> Set VName
forall als. TypeBase Size als -> Set VName
sizeNamesPos (StructType -> Set VName) -> StructType -> Set VName
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> ResRetType -> StructType
funType [Pat ParamType]
params' (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Size Uniqueness
ret
          hide :: VName -> (Int, b) -> Bool
hide VName
k (Int
lvl, b
_) =
            Int
lvl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
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

      Set VName
hidden_sizes <-
        [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName)
-> (Constraints -> [VName]) -> Constraints -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constraints -> [VName]
forall k a. Map k a -> [k]
M.keys (Constraints -> [VName])
-> (Constraints -> Constraints) -> Constraints -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> (Int, Constraint) -> Bool) -> Constraints -> Constraints
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey VName -> (Int, Constraint) -> Bool
forall {b}. VName -> (Int, b) -> Bool
hide (Constraints -> Set VName) -> m Constraints -> m (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints

      let onDim :: VName -> Set VName
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

      ResRetType -> m ResRetType
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResRetType -> m ResRetType) -> ResRetType -> m ResRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [VName]
forall a b. (a -> b) -> a -> b
$ (VName -> Set VName) -> Set VName -> Set VName
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> Set VName
onDim (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
$ TypeBase Size Uniqueness -> FV
forall u. TypeBase Size u -> FV
freeInType TypeBase Size Uniqueness
ret) TypeBase Size Uniqueness
ret
checkExp (OpSection QualName VName
op NoInfo StructType
_ SrcLoc
loc) = do
  StructType
ftype <- SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
op
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
OpSection QualName VName
op (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ftype) SrcLoc
loc
checkExp (OpSectionLeft QualName VName
op NoInfo StructType
_ ExpBase NoInfo VName
e (NoInfo (PName, ParamType, Maybe VName), NoInfo (PName, ParamType))
_ (NoInfo ResRetType, NoInfo [VName])
_ SrcLoc
loc) = do
  StructType
ftype <- SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
op
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  (StructType
t1, StructType
rt, Maybe VName
argext, [VName]
retext) <- SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
op, Int
0) StructType
ftype Size
e'
  case (StructType
ftype, StructType
rt) of
    (Scalar (Arrow NoUniqueness
_ PName
m1 Diet
d1 StructType
_ ResRetType
_), Scalar (Arrow NoUniqueness
_ PName
m2 Diet
d2 StructType
t2 ResRetType
rettype)) ->
      Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
        QualName VName
-> Info StructType
-> Size
-> (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
-> (Info ResRetType, Info [VName])
-> SrcLoc
-> Size
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)
          Size
e'
          ((PName, ParamType, Maybe VName)
-> Info (PName, ParamType, Maybe VName)
forall a. a -> Info a
Info (PName
m1, Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Size u -> 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 u. Diet -> TypeBase Size u -> ParamType
toParam Diet
d2 StructType
t2))
          (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
rettype, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
retext)
          SrcLoc
loc
    (StructType, StructType)
_ ->
      SrcLoc -> Notes -> Doc () -> TermTypeM Size
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 Size) -> Doc () -> TermTypeM Size
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 NoInfo StructType
_ ExpBase NoInfo VName
e (NoInfo (PName, ParamType), NoInfo (PName, ParamType, Maybe VName))
_ NoInfo ResRetType
NoInfo SrcLoc
loc) = do
  StructType
ftype <- SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
loc QualName VName
op
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  case StructType
ftype of
    Scalar (Arrow NoUniqueness
_ PName
m1 Diet
d1 StructType
t1 (RetType [] (Scalar (Arrow Uniqueness
_ PName
m2 Diet
d2 StructType
t2 (RetType [VName]
dims2 TypeBase Size Uniqueness
ret))))) -> do
      (StructType
t2', StructType
arrow', Maybe VName
argext, [VName]
_) <-
        SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply
          SrcLoc
loc
          (QualName VName -> Maybe (QualName VName)
forall a. a -> Maybe a
Just QualName VName
op, Int
1)
          (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size NoUniqueness -> StructType)
-> ScalarTypeBase Size NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
m2 Diet
d2 StructType
t2 (ResRetType -> ScalarTypeBase Size NoUniqueness)
-> ResRetType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase Size Uniqueness -> ResRetType)
-> TypeBase Size Uniqueness -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Size Uniqueness -> TypeBase Size Uniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Uniqueness -> TypeBase Size Uniqueness)
-> ScalarTypeBase Size Uniqueness -> TypeBase Size Uniqueness
forall a b. (a -> b) -> a -> b
$ Uniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Size Uniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow Uniqueness
Nonunique PName
m1 Diet
d1 StructType
t1 (ResRetType -> ScalarTypeBase Size Uniqueness)
-> ResRetType -> ScalarTypeBase Size Uniqueness
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims2 TypeBase Size Uniqueness
ret)
          Size
e'
      case StructType
arrow' of
        Scalar (Arrow NoUniqueness
_ PName
_ Diet
_ StructType
t1' (RetType [VName]
dims2' TypeBase Size Uniqueness
ret')) ->
          Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
            QualName VName
-> Info StructType
-> Size
-> (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
-> Info ResRetType
-> SrcLoc
-> Size
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)
              Size
e'
              ((PName, ParamType) -> Info (PName, ParamType)
forall a. a -> Info a
Info (PName
m1, Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Size u -> 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 u. Diet -> TypeBase Size u -> 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] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims2' TypeBase Size Uniqueness
ret')
              SrcLoc
loc
        StructType
_ -> [Char] -> TermTypeM Size
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermTypeM Size) -> [Char] -> TermTypeM Size
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 Size
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 Size) -> Doc () -> TermTypeM Size
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 (ProjectSection [Name]
fields NoInfo StructType
NoInfo SrcLoc
loc) = do
  StructType
a <- SrcLoc -> Name -> TermTypeM StructType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"a"
  let usage :: Usage
usage = SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"projection at"
  StructType
b <- (StructType -> Name -> TermTypeM StructType)
-> StructType -> [Name] -> TermTypeM StructType
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ((Name -> StructType -> TermTypeM StructType)
-> StructType -> Name -> TermTypeM StructType
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Name -> StructType -> TermTypeM StructType)
 -> StructType -> Name -> TermTypeM StructType)
-> (Name -> StructType -> TermTypeM StructType)
-> StructType
-> Name
-> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ Usage -> Name -> StructType -> TermTypeM StructType
forall (m :: * -> *).
MonadUnify m =>
Usage -> Name -> StructType -> m StructType
mustHaveField Usage
usage) StructType
a [Name]
fields
  let ft :: StructType
ft = ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size NoUniqueness -> StructType)
-> ScalarTypeBase Size NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe StructType
a (ResRetType -> ScalarTypeBase Size NoUniqueness)
-> ResRetType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase Size Uniqueness -> ResRetType)
-> TypeBase Size Uniqueness -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> TypeBase Size Uniqueness
forall u. Uniqueness -> TypeBase Size u -> TypeBase Size Uniqueness
toRes Uniqueness
Nonunique StructType
b
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [Name] -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
[Name] -> f StructType -> SrcLoc -> ExpBase f vn
ProjectSection [Name]
fields (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ft) SrcLoc
loc
checkExp (IndexSection SliceBase NoInfo VName
slice NoInfo StructType
NoInfo SrcLoc
loc) = do
  [DimIndex]
slice' <- SliceBase NoInfo VName -> TermTypeM [DimIndex]
checkSlice SliceBase NoInfo VName
slice
  (StructType
t, StructType
_) <- Usage -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType (SrcLoc -> Usage
forall a. Located a => a -> Usage
mkUsage' SrcLoc
loc) Name
"e" (Int -> TermTypeM (StructType, StructType))
-> Int -> TermTypeM (StructType, StructType)
forall a b. (a -> b) -> a -> b
$ [DimIndex] -> Int
sliceDims [DimIndex]
slice'
  (StructType
t', [VName]
retext) <- Maybe (SrcLoc, Rigidity)
-> [DimIndex] -> StructType -> TermTypeM (StructType, [VName])
forall as.
Maybe (SrcLoc, Rigidity)
-> [DimIndex]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
sliceShape Maybe (SrcLoc, Rigidity)
forall a. Maybe a
Nothing [DimIndex]
slice' StructType
t
  let ft :: StructType
ft = ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size NoUniqueness -> StructType)
-> ScalarTypeBase Size NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe StructType
t (ResRetType -> ScalarTypeBase Size NoUniqueness)
-> ResRetType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
retext (TypeBase Size Uniqueness -> ResRetType)
-> TypeBase Size Uniqueness -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> TypeBase Size Uniqueness
forall u. Uniqueness -> TypeBase Size u -> TypeBase Size Uniqueness
toRes Uniqueness
Nonunique StructType
t'
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ [DimIndex] -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
SliceBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
IndexSection [DimIndex]
slice' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ft) SrcLoc
loc
checkExp (AppExp (Loop [VName]
_ PatBase NoInfo VName ParamType
mergepat LoopInitBase NoInfo VName
loopinit LoopFormBase NoInfo VName
form ExpBase NoInfo VName
loopbody SrcLoc
loc) NoInfo AppRes
_) = do
  (([VName]
sparams, Pat ParamType
mergepat', LoopInitBase Info VName
loopinit', LoopFormBase Info VName
form', Size
loopbody'), AppRes
appres) <-
    (ExpBase NoInfo VName -> TermTypeM Size)
-> UncheckedLoop -> SrcLoc -> TermTypeM (CheckedLoop, AppRes)
checkLoop ExpBase NoInfo VName -> TermTypeM Size
checkExp (PatBase NoInfo VName ParamType
mergepat, LoopInitBase NoInfo VName
loopinit, LoopFormBase NoInfo VName
form, ExpBase NoInfo VName
loopbody) SrcLoc
loc
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
    AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
      ([VName]
-> Pat ParamType
-> LoopInitBase Info VName
-> LoopFormBase Info VName
-> Size
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[VName]
-> PatBase f vn ParamType
-> LoopInitBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
Loop [VName]
sparams Pat ParamType
mergepat' LoopInitBase Info VName
loopinit' LoopFormBase Info VName
form' Size
loopbody' SrcLoc
loc)
      (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
appres)
checkExp (Constr Name
name [ExpBase NoInfo VName]
es NoInfo StructType
NoInfo SrcLoc
loc) = do
  StructType
t <- SrcLoc -> Name -> TermTypeM StructType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
  [Size]
es' <- (ExpBase NoInfo VName -> TermTypeM Size)
-> [ExpBase NoInfo VName] -> TermTypeM [Size]
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 ExpBase NoInfo VName -> TermTypeM Size
checkExp [ExpBase NoInfo VName]
es
  [StructType]
ets <- (Size -> TermTypeM StructType) -> [Size] -> TermTypeM [StructType]
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 Size -> TermTypeM StructType
expType [Size]
es'
  Usage -> Name -> StructType -> [StructType] -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> Name -> StructType -> [StructType] -> m ()
mustHaveConstr (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"use of constructor") Name
name StructType
t [StructType]
ets
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ Name -> [Size] -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
Constr Name
name [Size]
es' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc
checkExp (AppExp (Match ExpBase NoInfo VName
e NonEmpty (CaseBase NoInfo VName)
cs SrcLoc
loc) NoInfo AppRes
_) = do
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  StructType
mt <- Size -> TermTypeM StructType
expType Size
e'
  (NonEmpty (CaseBase Info VName)
cs', StructType
t, [VName]
retext) <- StructType
-> NonEmpty (CaseBase NoInfo VName)
-> TermTypeM (NonEmpty (CaseBase Info VName), StructType, [VName])
checkCases StructType
mt NonEmpty (CaseBase NoInfo VName)
cs
  Usage -> Text -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> Text -> StructType -> m ()
zeroOrderType
    (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"being returned 'match'")
    Text
"type returned from pattern match"
    StructType
t
  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> TermTypeM Size) -> Size -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Size
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Size
-> NonEmpty (CaseBase Info VName)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match Size
e' NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
t [VName]
retext)
checkExp (Attr AttrInfo VName
info ExpBase NoInfo VName
e SrcLoc
loc) =
  AttrInfo VName -> Size -> SrcLoc -> Size
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr (AttrInfo VName -> Size -> SrcLoc -> Size)
-> TermTypeM (AttrInfo VName) -> TermTypeM (Size -> SrcLoc -> Size)
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 (Size -> SrcLoc -> Size)
-> TermTypeM Size -> TermTypeM (SrcLoc -> Size)
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e TermTypeM (SrcLoc -> Size) -> TermTypeM SrcLoc -> TermTypeM Size
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

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

checkCase ::
  StructType ->
  CaseBase NoInfo VName ->
  TermTypeM (CaseBase Info VName, StructType, [VName])
checkCase :: StructType
-> CaseBase NoInfo VName
-> TermTypeM (CaseBase Info VName, StructType, [VName])
checkCase StructType
mt (CasePat PatBase NoInfo VName StructType
p ExpBase NoInfo VName
e SrcLoc
loc) =
  [SizeBinder VName]
-> PatBase NoInfo VName StructType
-> StructType
-> (Pat ParamType
    -> TermTypeM (CaseBase Info VName, StructType, [VName]))
-> TermTypeM (CaseBase Info VName, StructType, [VName])
forall u a.
[SizeBinder VName]
-> PatBase NoInfo VName (TypeBase Size u)
-> StructType
-> (Pat ParamType -> TermTypeM a)
-> TermTypeM a
bindingPat [] PatBase NoInfo 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
    Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
    StructType
e_t <- Size -> TermTypeM StructType
expTypeFully Size
e'
    (StructType
e_t', [VName]
retext) <- SrcLoc -> [VName] -> StructType -> TermTypeM (StructType, [VName])
forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType SrcLoc
loc (Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames Pat ParamType
p') StructType
e_t
    (CaseBase Info VName, StructType, [VName])
-> TermTypeM (CaseBase Info VName, StructType, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatBase Info VName StructType
-> Size -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat ((ParamType -> StructType)
-> Pat ParamType -> PatBase Info VName StructType
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct Pat ParamType
p') Size
e' SrcLoc
loc, StructType
e_t', [VName]
retext)

-- | An unmatched pattern. Used in in the generation of
-- unmatched pattern warnings by the type checker.
data Unmatched p
  = UnmatchedNum p [PatLit]
  | UnmatchedBool p
  | UnmatchedConstr p
  | Unmatched p
  deriving ((forall a b. (a -> b) -> Unmatched a -> Unmatched b)
-> (forall a b. a -> Unmatched b -> Unmatched a)
-> Functor Unmatched
forall a b. a -> Unmatched b -> Unmatched a
forall a b. (a -> b) -> Unmatched a -> Unmatched b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Unmatched a -> Unmatched b
fmap :: forall a b. (a -> b) -> Unmatched a -> Unmatched b
$c<$ :: forall a b. a -> Unmatched b -> Unmatched a
<$ :: forall a b. a -> Unmatched b -> Unmatched a
Functor, Int -> Unmatched p -> [Char] -> [Char]
[Unmatched p] -> [Char] -> [Char]
Unmatched p -> [Char]
(Int -> Unmatched p -> [Char] -> [Char])
-> (Unmatched p -> [Char])
-> ([Unmatched p] -> [Char] -> [Char])
-> Show (Unmatched p)
forall p. Show p => Int -> Unmatched p -> [Char] -> [Char]
forall p. Show p => [Unmatched p] -> [Char] -> [Char]
forall p. Show p => Unmatched p -> [Char]
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: forall p. Show p => Int -> Unmatched p -> [Char] -> [Char]
showsPrec :: Int -> Unmatched p -> [Char] -> [Char]
$cshow :: forall p. Show p => Unmatched p -> [Char]
show :: Unmatched p -> [Char]
$cshowList :: forall p. Show p => [Unmatched p] -> [Char] -> [Char]
showList :: [Unmatched p] -> [Char] -> [Char]
Show)

instance Pretty (Unmatched (Pat StructType)) where
  pretty :: forall ann. Unmatched (PatBase Info VName StructType) -> Doc ann
pretty Unmatched (PatBase Info VName StructType)
um = case Unmatched (PatBase Info VName StructType)
um of
    (UnmatchedNum PatBase Info VName StructType
p [PatLit]
nums) -> PatBase Info VName StructType -> Doc ann
forall {v} {f :: * -> *} {t} {ann}.
(Eq v, IsName v, Annot f, Pretty t) =>
PatBase f v t -> Doc ann
pretty' PatBase Info VName StructType
p Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
"where p is not one of" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [PatLit] -> Doc ann
forall ann. [PatLit] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [PatLit]
nums
    (UnmatchedBool PatBase Info VName StructType
p) -> PatBase Info VName StructType -> Doc ann
forall {v} {f :: * -> *} {t} {ann}.
(Eq v, IsName v, Annot f, Pretty t) =>
PatBase f v t -> Doc ann
pretty' PatBase Info VName StructType
p
    (UnmatchedConstr PatBase Info VName StructType
p) -> PatBase Info VName StructType -> Doc ann
forall {v} {f :: * -> *} {t} {ann}.
(Eq v, IsName v, Annot f, Pretty t) =>
PatBase f v t -> Doc ann
pretty' PatBase Info VName StructType
p
    (Unmatched PatBase Info VName StructType
p) -> PatBase Info VName StructType -> Doc ann
forall {v} {f :: * -> *} {t} {ann}.
(Eq v, IsName v, Annot f, Pretty t) =>
PatBase f v t -> Doc ann
pretty' PatBase Info VName StructType
p
    where
      pretty' :: PatBase f v t -> Doc ann
pretty' (PatAscription PatBase f v t
p TypeExp (ExpBase f v) v
t SrcLoc
_) = PatBase f v t -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. PatBase f v t -> Doc ann
pretty PatBase f v t
p Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
":" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> TypeExp (ExpBase f v) v -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeExp (ExpBase f v) v -> Doc ann
pretty TypeExp (ExpBase f v) v
t
      pretty' (PatParens PatBase f v t
p SrcLoc
_) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ PatBase f v t -> Doc ann
pretty' PatBase f v t
p
      pretty' (PatAttr AttrInfo v
_ PatBase f v t
p SrcLoc
_) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ PatBase f v t -> Doc ann
pretty' PatBase f v t
p
      pretty' (Id v
v f t
_ SrcLoc
_) = v -> Doc ann
forall a. v -> Doc a
forall v a. IsName v => v -> Doc a
prettyName v
v
      pretty' (TuplePat [PatBase f v t]
pats SrcLoc
_) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (PatBase f v t -> Doc ann) -> [PatBase f v t] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map PatBase f v t -> Doc ann
pretty' [PatBase f v t]
pats
      pretty' (RecordPat [(L Name, PatBase f v t)]
fs SrcLoc
_) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
braces (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ ((L Name, PatBase f v t) -> Doc ann)
-> [(L Name, PatBase f v t)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, PatBase f v t) -> Doc ann
ppField [(L Name, PatBase f v t)]
fs
        where
          ppField :: (L Name, PatBase f v t) -> Doc ann
ppField (L Loc
_ Name
name, PatBase f v t
t) = [Char] -> Doc ann
forall ann. [Char] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Name -> [Char]
nameToString Name
name) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
equals Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> PatBase f v t -> Doc ann
pretty' PatBase f v t
t
      pretty' Wildcard {} = Doc ann
"_"
      pretty' (PatLit PatLit
e f t
_ SrcLoc
_) = PatLit -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. PatLit -> Doc ann
pretty PatLit
e
      pretty' (PatConstr Name
n f t
_ [PatBase f v t]
ps SrcLoc
_) = Doc ann
"#" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Name -> Doc ann
pretty Name
n Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
sep ((PatBase f v t -> Doc ann) -> [PatBase f v t] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map PatBase f v t -> Doc ann
pretty' [PatBase f v t]
ps)

checkIdent :: IdentBase NoInfo VName StructType -> TermTypeM (Ident StructType)
checkIdent :: IdentBase NoInfo VName StructType -> TermTypeM (Ident StructType)
checkIdent (Ident VName
name NoInfo StructType
_ SrcLoc
loc) = do
  StructType
vt <- SrcLoc -> QualName VName -> TermTypeM StructType
lookupVar SrcLoc
loc (QualName VName -> TermTypeM StructType)
-> QualName VName -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
name
  Ident StructType -> TermTypeM (Ident StructType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ident StructType -> TermTypeM (Ident StructType))
-> Ident StructType -> TermTypeM (Ident StructType)
forall a b. (a -> b) -> a -> b
$ VName -> Info StructType -> SrcLoc -> Ident StructType
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident VName
name (StructType -> Info StructType
forall a. a -> Info a
Info StructType
vt) SrcLoc
loc

checkSlice :: SliceBase NoInfo VName -> TermTypeM [DimIndex]
checkSlice :: SliceBase NoInfo VName -> TermTypeM [DimIndex]
checkSlice = (DimIndexBase NoInfo VName -> TermTypeM DimIndex)
-> SliceBase NoInfo VName -> 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 DimIndexBase NoInfo VName -> TermTypeM DimIndex
checkDimIndex
  where
    checkDimIndex :: DimIndexBase NoInfo VName -> TermTypeM DimIndex
checkDimIndex (DimFix ExpBase NoInfo VName
i) = do
      Size -> DimIndex
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (Size -> DimIndex) -> TermTypeM Size -> TermTypeM DimIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> [PrimType] -> Size -> TermTypeM Size
require Text
"use as index" [PrimType]
anySignedType (Size -> TermTypeM Size) -> TermTypeM Size -> TermTypeM Size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
i)
    checkDimIndex (DimSlice Maybe (ExpBase NoInfo VName)
i Maybe (ExpBase NoInfo VName)
j Maybe (ExpBase NoInfo VName)
s) =
      Maybe Size -> Maybe Size -> Maybe Size -> DimIndex
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe Size -> Maybe Size -> Maybe Size -> DimIndex)
-> TermTypeM (Maybe Size)
-> TermTypeM (Maybe Size -> Maybe Size -> DimIndex)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (ExpBase NoInfo VName) -> TermTypeM (Maybe Size)
check Maybe (ExpBase NoInfo VName)
i TermTypeM (Maybe Size -> Maybe Size -> DimIndex)
-> TermTypeM (Maybe Size) -> TermTypeM (Maybe Size -> DimIndex)
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (ExpBase NoInfo VName) -> TermTypeM (Maybe Size)
check Maybe (ExpBase NoInfo VName)
j TermTypeM (Maybe Size -> DimIndex)
-> TermTypeM (Maybe Size) -> 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
<*> Maybe (ExpBase NoInfo VName) -> TermTypeM (Maybe Size)
check Maybe (ExpBase NoInfo VName)
s

    check :: Maybe (ExpBase NoInfo VName) -> TermTypeM (Maybe Size)
check =
      TermTypeM (Maybe Size)
-> (ExpBase NoInfo VName -> TermTypeM (Maybe Size))
-> Maybe (ExpBase NoInfo VName)
-> TermTypeM (Maybe Size)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe Size -> TermTypeM (Maybe Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Size
forall a. Maybe a
Nothing) ((ExpBase NoInfo VName -> TermTypeM (Maybe Size))
 -> Maybe (ExpBase NoInfo VName) -> TermTypeM (Maybe Size))
-> (ExpBase NoInfo VName -> TermTypeM (Maybe Size))
-> Maybe (ExpBase NoInfo VName)
-> TermTypeM (Maybe Size)
forall a b. (a -> b) -> a -> b
$
        (Size -> Maybe Size) -> TermTypeM Size -> TermTypeM (Maybe Size)
forall a b. (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Size -> Maybe Size
forall a. a -> Maybe a
Just (TermTypeM Size -> TermTypeM (Maybe Size))
-> (Size -> TermTypeM Size) -> Size -> TermTypeM (Maybe Size)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> StructType -> Size -> TermTypeM Size
unifies Text
"use as index" (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size NoUniqueness -> StructType)
-> ScalarTypeBase Size NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Size NoUniqueness)
-> PrimType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64) (Size -> TermTypeM (Maybe Size))
-> (ExpBase NoInfo VName -> TermTypeM Size)
-> ExpBase NoInfo VName
-> TermTypeM (Maybe Size)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< ExpBase NoInfo VName -> TermTypeM Size
checkExp

-- The number of dimensions affected by this slice (so the minimum
-- rank of the array we are slicing).
sliceDims :: [DimIndex] -> Int
sliceDims :: [DimIndex] -> Int
sliceDims = [DimIndex] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

instantiateDimsInReturnType ::
  SrcLoc ->
  Maybe (QualName VName) ->
  ResRetType ->
  TermTypeM (ResType, [VName])
instantiateDimsInReturnType :: SrcLoc
-> Maybe (QualName VName)
-> ResRetType
-> TermTypeM (TypeBase Size Uniqueness, [VName])
instantiateDimsInReturnType SrcLoc
loc Maybe (QualName VName)
fname (RetType [VName]
dims TypeBase Size Uniqueness
t)
  | [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
dims =
      (TypeBase Size Uniqueness, [VName])
-> TermTypeM (TypeBase Size Uniqueness, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Size Uniqueness
t, [VName]
forall a. Monoid a => a
mempty)
  | Bool
otherwise = do
      [VName]
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
      (TypeBase Size Uniqueness, [VName])
-> TermTypeM (TypeBase Size Uniqueness, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Size -> Size)
-> TypeBase Size Uniqueness -> TypeBase Size Uniqueness
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ([(VName, Subst StructRetType)] -> Size -> Size
forall {a}.
Substitutable a =>
[(VName, Subst StructRetType)] -> a -> a
onDim ([(VName, Subst StructRetType)] -> Size -> Size)
-> [(VName, Subst StructRetType)] -> Size -> Size
forall a b. (a -> b) -> a -> b
$ [VName] -> [Subst StructRetType] -> [(VName, Subst StructRetType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
dims ([Subst StructRetType] -> [(VName, Subst StructRetType)])
-> [Subst StructRetType] -> [(VName, Subst StructRetType)]
forall a b. (a -> b) -> a -> b
$ (VName -> Subst StructRetType) -> [VName] -> [Subst StructRetType]
forall a b. (a -> b) -> [a] -> [b]
map (Size -> Subst StructRetType
forall t. Size -> Subst t
ExpSubst (Size -> Subst StructRetType)
-> (VName -> Size) -> VName -> Subst StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> SrcLoc -> Size
`sizeFromName` SrcLoc
loc) (QualName VName -> Size)
-> (VName -> QualName VName) -> VName -> Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName) [VName]
dims') TypeBase Size Uniqueness
t, [VName]
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
. [Char] -> Name
nameFromString
        ([Char] -> Name) -> (VName -> [Char]) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> [Char] -> [Char]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Char -> Bool
isAscii
        ([Char] -> [Char]) -> (VName -> [Char]) -> VName -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
baseString
    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 as -> S.Set VName
boundInsideType :: forall als. TypeBase Size als -> Set VName
boundInsideType (Array as
_ Shape Size
_ ScalarTypeBase Size NoUniqueness
t) = StructType -> Set VName
forall als. TypeBase Size als -> Set VName
boundInsideType (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase Size NoUniqueness
t)
boundInsideType (Scalar Prim {}) = Set VName
forall a. Monoid a => a
mempty
boundInsideType (Scalar (TypeVar as
_ QualName VName
_ [TypeArg Size]
targs)) = (TypeArg Size -> Set VName) -> [TypeArg Size] -> 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 Size -> Set VName
f [TypeArg Size]
targs
  where
    f :: TypeArg Size -> Set VName
f (TypeArgType StructType
t) = StructType -> Set VName
forall als. TypeBase Size als -> Set VName
boundInsideType StructType
t
    f TypeArgDim {} = Set VName
forall a. Monoid a => a
mempty
boundInsideType (Scalar (Record Map Name (TypeBase Size as)
fs)) = (TypeBase Size as -> Set VName)
-> Map Name (TypeBase Size as) -> 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 Size as -> Set VName
forall als. TypeBase Size als -> Set VName
boundInsideType Map Name (TypeBase Size as)
fs
boundInsideType (Scalar (Sum Map Name [TypeBase Size as]
cs)) = ([TypeBase Size as] -> Set VName)
-> Map Name [TypeBase Size as] -> 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 Size as -> Set VName) -> [TypeBase Size as] -> 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 Size as -> Set VName
forall als. TypeBase Size als -> Set VName
boundInsideType) Map Name [TypeBase Size as]
cs
boundInsideType (Scalar (Arrow as
_ PName
pn Diet
_ StructType
t1 (RetType [VName]
dims TypeBase Size Uniqueness
t2))) =
  Set VName
pn' Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> StructType -> Set VName
forall als. TypeBase Size als -> 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
<> TypeBase Size Uniqueness -> Set VName
forall als. TypeBase Size als -> Set VName
boundInsideType TypeBase Size Uniqueness
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 u -> (Names, Names)
dimUses :: forall u. TypeBase Size u -> (Set VName, Set VName)
dimUses = (State (Set VName, Set VName) (TypeBase () u)
 -> (Set VName, Set VName) -> (Set VName, Set VName))
-> (Set VName, Set VName)
-> State (Set VName, Set VName) (TypeBase () u)
-> (Set VName, Set VName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (Set VName, Set VName) (TypeBase () u)
-> (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 () u)
 -> (Set VName, Set VName))
-> (TypeBase Size u
    -> State (Set VName, Set VName) (TypeBase () u))
-> TypeBase Size u
-> (Set VName, Set VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName
 -> DimPos -> Size -> StateT (Set VName, Set VName) Identity ())
-> TypeBase Size u -> State (Set VName, Set VName) (TypeBase () u)
forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName
-> DimPos -> Size -> StateT (Set VName, Set VName) Identity ()
forall {m :: * -> *}.
MonadState (Set VName, Set VName) m =>
Set VName -> DimPos -> Size -> m ()
f
  where
    f :: Set VName -> DimPos -> Size -> m ()
f Set VName
bound DimPos
pos Size
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) <>)
        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) <>)
        DimPos
PosReturn -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      where
        fv :: FV
fv = Size -> FV
freeInExp Size
e FV -> Set VName -> FV
`freeWithout` Set VName
bound

checkApply ::
  SrcLoc ->
  ApplyOp ->
  StructType ->
  Exp ->
  TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply :: SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc (Maybe (QualName VName)
fname, Int
_) (Scalar (Arrow NoUniqueness
_ PName
pname Diet
_ StructType
tp1 ResRetType
tp2)) Size
argexp = do
  let argtype :: StructType
argtype = Size -> StructType
typeOf Size
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)
-> Size -> StructType -> StructType -> Checking
CheckingApply Maybe (QualName VName)
fname Size
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 (Size -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Size
argexp Text
"use as function argument") StructType
tp1 StructType
argtype

    -- Perform substitutions of instantiated variables in the types.
    (TypeBase Size Uniqueness
tp2', [VName]
ext) <- SrcLoc
-> Maybe (QualName VName)
-> ResRetType
-> TermTypeM (TypeBase Size Uniqueness, [VName])
instantiateDimsInReturnType SrcLoc
loc Maybe (QualName VName)
fname (ResRetType -> TermTypeM (TypeBase Size Uniqueness, [VName]))
-> TermTypeM ResRetType
-> TermTypeM (TypeBase Size Uniqueness, [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
    StructType
argtype' <- StructType -> TermTypeM StructType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully StructType
argtype

    -- Check whether this would produce an impossible return type.
    let (Set VName
tp2_produced_dims, Set VName
tp2_paramdims) = TypeBase Size Uniqueness -> (Set VName, Set VName)
forall u. TypeBase Size u -> (Set VName, Set VName)
dimUses TypeBase Size Uniqueness
tp2'
        problematic :: Set VName
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 als. TypeBase Size als -> Set VName
boundInsideType StructType
argtype'
        problem :: Bool
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)
    Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (Set VName -> Bool
forall a. Set a -> Bool
S.null Set VName
problematic) Bool -> Bool -> Bool
&& Bool
problem) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ do
      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 () -> Doc ()) -> Doc () -> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"existential-param-ret" (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Existential size would appear in function parameter of return type:"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ResRetType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ResRetType -> Doc ann
pretty ([VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext TypeBase Size Uniqueness
tp2'))
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
"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."

    (Maybe VName
argext, StructType
tp2'') <-
      case PName
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
$ TypeBase Size Uniqueness -> FV
forall u. TypeBase Size u -> FV
freeInType TypeBase Size Uniqueness
tp2') ->
              if Size -> Bool
hasBinding Size
argexp
                then do
                  SrcLoc -> Doc () -> TermTypeM ()
forall loc. Located loc => loc -> Doc () -> TermTypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn (Size -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Size
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."
                  VName
d <- Size -> 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 Size
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
$ Size -> ExpBase NoInfo VName
bareExp Size
argexp) Name
"n"
                  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
$ Size -> Subst t
forall t. Size -> Subst t
ExpSubst (Size -> Subst t) -> Size -> Subst t
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Size
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
d) (SrcLoc -> Size) -> SrcLoc -> Size
forall a b. (a -> b) -> a -> b
$ Size -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Size
argexp
                          else Maybe (Subst t)
forall a. Maybe a
Nothing
                  (Maybe VName, StructType) -> TermTypeM (Maybe VName, StructType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Maybe VName
forall a. a -> Maybe a
Just VName
d, 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
$ TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
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
$ Size -> Subst t
forall t. Size -> Subst t
ExpSubst (Size -> Subst t) -> Size -> Subst t
forall a b. (a -> b) -> a -> b
$ Size -> Maybe Size -> Size
forall a. a -> Maybe a -> a
fromMaybe Size
argexp (Maybe Size -> Size) -> Maybe Size -> Size
forall a b. (a -> b) -> a -> b
$ Size -> Maybe Size
stripExp Size
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
$ TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
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, TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
tp2')

    (StructType, StructType, Maybe VName, [VName])
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType
tp1, StructType
tp2'', Maybe VName
argext, [VName]
ext)
checkApply SrcLoc
loc ApplyOp
fname tfun :: StructType
tfun@(Scalar TypeVar {}) Size
arg = do
  ParamType
tv <- SrcLoc -> Name -> TermTypeM ParamType
forall als a dim.
(Monoid als, Located a) =>
a -> Name -> TermTypeM (TypeBase dim als)
forall (m :: * -> *) als a dim.
(MonadUnify m, Monoid als, Located a) =>
a -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"b"
  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
"use as function") StructType
tfun (StructType -> TermTypeM ()) -> StructType -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe (Size -> StructType
typeOf Size
arg) (ResRetType -> ScalarTypeBase Size NoUniqueness)
-> ResRetType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase Size Uniqueness -> ResRetType)
-> TypeBase Size Uniqueness -> ResRetType
forall a b. (a -> b) -> a -> b
$ ParamType -> TypeBase Size Uniqueness
paramToRes ParamType
tv)
  StructType
tfun' <- StructType -> TermTypeM StructType
forall (m :: * -> *). MonadUnify m => StructType -> m StructType
normType StructType
tfun
  SrcLoc
-> ApplyOp
-> StructType
-> Size
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
checkApply SrcLoc
loc ApplyOp
fname StructType
tfun' Size
arg
checkApply SrcLoc
loc (Maybe (QualName VName)
fname, Int
prev_applied) StructType
ftype Size
argexp = do
  let fname' :: Doc ann
fname' = Doc ann
-> (QualName VName -> Doc ann) -> Maybe (QualName VName) -> Doc ann
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Doc ann
"expression" (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
dquotes (Doc ann -> Doc ann)
-> (QualName VName -> Doc ann) -> QualName VName -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QualName VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName VName -> Doc ann
pretty) Maybe (QualName VName)
fname

  SrcLoc
-> Notes
-> Doc ()
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
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 (StructType, StructType, Maybe VName, [VName]))
-> Doc ()
-> TermTypeM (StructType, StructType, Maybe VName, [VName])
forall a b. (a -> b) -> a -> b
$
    if Int
prev_applied Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
      then
        Doc ()
"Cannot apply"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
forall ann. Doc ann
fname'
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"as function, as it has type:"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (StructType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty StructType
ftype)
      else
        Doc ()
"Cannot apply"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
forall ann. Doc ann
fname'
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"to argument #"
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ()
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Int
prev_applied Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Doc Any -> Doc ()
forall a b. Doc a -> Doc b
shorten (Doc Any -> Doc ()) -> Doc Any -> Doc ()
forall a b. (a -> b) -> a -> b
$ Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann
group (Doc Any -> Doc Any) -> Doc Any -> Doc Any
forall a b. (a -> b) -> a -> b
$ Size -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. Size -> Doc ann
pretty Size
argexp)
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
","
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"as"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
forall ann. Doc ann
fname'
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"only takes"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Int -> Doc ()
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Int
prev_applied
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
arguments
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
  where
    arguments :: Doc ()
arguments
      | Int
prev_applied Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Doc ()
"argument"
      | Bool
otherwise = Doc ()
"arguments"

-- | 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], Size)
checkOneExp ExpBase NoInfo VName
e = (ExpBase NoInfo VName -> TermTypeM Size)
-> TermTypeM ([TypeParamBase VName], Size)
-> TypeM ([TypeParamBase VName], Size)
forall a.
(ExpBase NoInfo VName -> TermTypeM Size) -> TermTypeM a -> TypeM a
runTermTypeM ExpBase NoInfo VName -> TermTypeM Size
checkExp (TermTypeM ([TypeParamBase VName], Size)
 -> TypeM ([TypeParamBase VName], Size))
-> TermTypeM ([TypeParamBase VName], Size)
-> TypeM ([TypeParamBase VName], Size)
forall a b. (a -> b) -> a -> b
$ do
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  let t :: StructType
t = Size -> StructType
typeOf Size
e'
  ([TypeParamBase VName]
tparams, [Pat ParamType]
_, ResRetType
_) <-
    Name
-> SrcLoc
-> [TypeParamBase VName]
-> [Pat ParamType]
-> TypeBase Size Uniqueness
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
letGeneralise ([Char] -> Name
nameFromString [Char]
"<exp>") (ExpBase NoInfo VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase NoInfo VName
e) [] [] (TypeBase Size Uniqueness
 -> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType))
-> TypeBase Size Uniqueness
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> TypeBase Size Uniqueness
forall u. Uniqueness -> TypeBase Size u -> TypeBase Size Uniqueness
toRes Uniqueness
Nonunique StructType
t
  Set VName -> TermTypeM ()
fixOverloadedTypes (Set VName -> TermTypeM ()) -> Set VName -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ StructType -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars StructType
t
  Size
e'' <- Size -> TermTypeM Size
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully Size
e'
  Size -> TermTypeM ()
localChecks Size
e''
  Size -> TermTypeM ()
causalityCheck Size
e''
  ([TypeParamBase VName], Size)
-> TermTypeM ([TypeParamBase VName], Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TypeParamBase VName]
tparams, Size
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 Size
checkSizeExp ExpBase NoInfo VName
e = (ExpBase NoInfo VName -> TermTypeM Size)
-> TermTypeM Size -> TypeM Size
forall a.
(ExpBase NoInfo VName -> TermTypeM Size) -> TermTypeM a -> TypeM a
runTermTypeM ExpBase NoInfo VName -> TermTypeM Size
checkExp (TermTypeM Size -> TypeM Size) -> TermTypeM Size -> TypeM Size
forall a b. (a -> b) -> a -> b
$ do
  Size
e' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
e
  let t :: StructType
t = Size -> StructType
typeOf Size
e'
  Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Size -> Bool
hasBinding Size
e') (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    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 (Size -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Size
e') Notes
forall a. Monoid a => a
mempty (Doc () -> TermTypeM ())
-> (Doc () -> Doc ()) -> Doc () -> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"size-expression-bind" (Doc () -> TermTypeM ()) -> Doc () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
      Doc ()
"Size expression with binding is forbidden."
  Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (Size -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Size
e' Text
"Size expression") StructType
t (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Size NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (IntType -> PrimType
Signed IntType
Int64)))
  Size -> TermTypeM Size
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully Size
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 :: Size -> TermTypeM ()
causalityCheck Size
binding_body = do
  Constraints
constraints <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints

  let checkCausality :: Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
what Set VName
known TypeBase Size u
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 Size u -> FV
forall u. TypeBase Size u -> FV
freeInType TypeBase Size u
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 Size u -> 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 Size u
t
        | Bool
otherwise = Maybe (t (Either TypeError) a)
forall a. Maybe a
Nothing

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

      collectingNewKnown :: StateT (Set VName) (Either TypeError) a
-> StateT (Set VName) (Either TypeError) (Set VName)
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 -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known (Var QualName VName
v (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> 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) Size
bad
      onExp Set VName
known (ProjectSection [Name]
_ (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"projection section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Size
bad
      onExp Set VName
known (IndexSection [DimIndex]
_ (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"projection section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Size
bad
      onExp Set VName
known (OpSectionRight QualName VName
_ (Info StructType
t) Size
_ (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
_ Info ResRetType
_ SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"operator section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Size
bad
      onExp Set VName
known (OpSectionLeft QualName VName
_ (Info StructType
t) Size
_ (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
_ (Info ResRetType, Info [VName])
_ SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"operator section" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Size
bad
      onExp Set VName
known (ArrayLit [] (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"empty array" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Size
bad
      onExp Set VName
known (Hole (Info StructType
t) SrcLoc
loc)
        | Just StateT (Set VName) (Either TypeError) Size
bad <- Doc ()
-> Set VName
-> StructType
-> SrcLoc
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> a
-> Maybe (t (Either TypeError) a)
checkCausality Doc ()
"hole" Set VName
known StructType
t SrcLoc
loc =
            StateT (Set VName) (Either TypeError) Size
bad
      onExp Set VName
known e :: Size
e@(Lambda [Pat ParamType]
params Size
body Maybe (TypeExp Size VName)
_ Info ResRetType
_ SrcLoc
_)
        | StateT (Set VName) (Either TypeError) Size
bad : [StateT (Set VName) (Either TypeError) Size]
_ <- (Pat ParamType
 -> Maybe (StateT (Set VName) (Either TypeError) Size))
-> [Pat ParamType] -> [StateT (Set VName) (Either TypeError) Size]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Set VName
-> Pat ParamType
-> Maybe (StateT (Set VName) (Either TypeError) Size)
forall {t :: (* -> *) -> * -> *} {u} {a}.
(MonadTrans t, Pretty u) =>
Set VName
-> Pat (TypeBase Size u) -> Maybe (t (Either TypeError) a)
checkParamCausality Set VName
known) [Pat ParamType]
params =
            StateT (Set VName) (Either TypeError) Size
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) Size
-> 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) Size
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known Size
body
            Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
      onExp Set VName
known e :: Size
e@(AppExp (LetPat [SizeBinder VName]
_ PatBase Info VName StructType
_ Size
bindee_e Size
body_e SrcLoc
_) (Info AppRes
res)) = do
        Set VName
-> Size
-> Size
-> [VName]
-> StateT (Set VName) (Either TypeError) ()
sequencePoint Set VName
known Size
bindee_e Size
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
        Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
      onExp Set VName
known e :: Size
e@(AppExp (Match Size
scrutinee NonEmpty (CaseBase Info VName)
cs SrcLoc
_) (Info AppRes
res)) = do
        Set VName
new_known <- StateT (Set VName) (Either TypeError) Size
-> 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) Size
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known Size
scrutinee
        StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) ()
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName
-> NonEmpty (CaseBase Info VName)
-> StateT (Set VName) (Either TypeError) ()
forall {a}.
ASTMappable a =>
Set VName -> a -> StateT (Set VName) (Either TypeError) ()
recurse (Set VName
new_known Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
known) NonEmpty (CaseBase Info VName)
cs
        (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName
new_known 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)) <>)
        Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
      onExp Set VName
known e :: Size
e@(AppExp (Apply Size
f NonEmpty (Info (Maybe VName), Size)
args SrcLoc
_) (Info AppRes
res)) = do
        Set VName
-> [(Info (Maybe VName), Size)]
-> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known ([(Info (Maybe VName), Size)]
 -> StateT (Set VName) (Either TypeError) ())
-> [(Info (Maybe VName), Size)]
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ [(Info (Maybe VName), Size)] -> [(Info (Maybe VName), Size)]
forall a. [a] -> [a]
reverse ([(Info (Maybe VName), Size)] -> [(Info (Maybe VName), Size)])
-> [(Info (Maybe VName), Size)] -> [(Info (Maybe VName), Size)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Size) -> [(Info (Maybe VName), Size)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Size)
args
        Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
        where
          seqArgs :: Set VName
-> [(Info (Maybe VName), Size)]
-> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known' [] = do
            StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) Size
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known' Size
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) <>)
          seqArgs Set VName
known' ((Info Maybe VName
p, Size
x) : [(Info (Maybe VName), Size)]
xs) = do
            Set VName
new_known <- StateT (Set VName) (Either TypeError) Size
-> 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) Size
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known' Size
x
            StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) ()
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName
-> [(Info (Maybe VName), Size)]
-> StateT (Set VName) (Either TypeError) ()
seqArgs (Set VName
new_known Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
known') [(Info (Maybe VName), Size)]
xs
            (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName
new_known 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 (Maybe VName -> [VName]
forall a. Maybe a -> [a]
maybeToList Maybe VName
p)) <>)
      onExp Set VName
known e :: Size
e@(Constr Name
v [Size]
args (Info StructType
t) SrcLoc
loc) = do
        Set VName -> [Size] -> StateT (Set VName) (Either TypeError) ()
seqArgs Set VName
known [Size]
args
        Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
        where
          seqArgs :: Set VName -> [Size] -> 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 :: (* -> *) -> * -> *} {u} {a} {a}.
(MonadTrans t, Pretty u, Located a) =>
Doc ()
-> Set VName
-> TypeBase Size u
-> 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' (Size
x : [Size]
xs) = do
            Set VName
new_known <- StateT (Set VName) (Either TypeError) Size
-> 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) Size
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known' Size
x
            StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) ()
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) ()
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName -> [Size] -> StateT (Set VName) (Either TypeError) ()
seqArgs (Set VName
new_known Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
known') [Size]
xs
            (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Set VName
new_known <>)
      onExp
        Set VName
known
        e :: Size
e@(AppExp (BinOp (QualName VName
f, SrcLoc
floc) Info StructType
ft (Size
x, Info Maybe VName
xp) (Size
y, Info Maybe VName
yp) SrcLoc
_) (Info AppRes
res)) = do
          Set VName
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
-> Size
-> Size
-> [VName]
-> StateT (Set VName) (Either TypeError) ()
sequencePoint Set VName
known Size
x Size
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]
          StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) Size
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp (Set VName
args_known Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
known) (QualName VName -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
f Info StructType
ft SrcLoc
floc)
          (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName
args_known 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)) <>)
          Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
      onExp Set VName
known e :: Size
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))
        Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e
      onExp Set VName
known Size
e = do
        Set VName -> Size -> StateT (Set VName) (Either TypeError) ()
forall {a}.
ASTMappable a =>
Set VName -> a -> StateT (Set VName) (Either TypeError) ()
recurse Set VName
known Size
e
        Size -> StateT (Set VName) (Either TypeError) Size
forall a. a -> StateT (Set VName) (Either TypeError) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
e

      recurse :: Set VName -> a -> StateT (Set VName) (Either TypeError) ()
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
-> Size
-> Size
-> [VName]
-> StateT (Set VName) (Either TypeError) ()
sequencePoint Set VName
known Size
x Size
y [VName]
ext = do
        Set VName
new_known <- StateT (Set VName) (Either TypeError) Size
-> 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) Size
 -> StateT (Set VName) (Either TypeError) (Set VName))
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
known Size
x
        StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set VName) (Either TypeError) Size
 -> StateT (Set VName) (Either TypeError) ())
-> StateT (Set VName) (Either TypeError) Size
-> StateT (Set VName) (Either TypeError) ()
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp (Set VName
new_known Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
known) Size
y
        (Set VName -> Set VName)
-> StateT (Set VName) (Either TypeError) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName
new_known 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]
ext) <>)

  (TypeError -> TermTypeM ())
-> (Size -> TermTypeM ()) -> Either TypeError Size -> TermTypeM ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either TypeError -> TermTypeM ()
forall a. TypeError -> TermTypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TermTypeM () -> Size -> TermTypeM ()
forall a b. a -> b -> a
const (TermTypeM () -> Size -> TermTypeM ())
-> TermTypeM () -> Size -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (Either TypeError Size -> TermTypeM ())
-> Either TypeError Size -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    StateT (Set VName) (Either TypeError) Size
-> Set VName -> Either TypeError Size
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (Set VName -> Size -> StateT (Set VName) (Either TypeError) Size
onExp Set VName
forall a. Monoid a => a
mempty Size
binding_body) Set VName
forall a. Monoid a => a
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
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
      (a, Loc) -> Maybe (a, Loc)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
v, Loc
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
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
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
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
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'))

-- | 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)
localChecks :: Exp -> TermTypeM ()
localChecks :: Size -> TermTypeM ()
localChecks = TermTypeM Size -> TermTypeM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TermTypeM Size -> TermTypeM ())
-> (Size -> TermTypeM Size) -> Size -> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Size -> TermTypeM Size
check
  where
    check :: Size -> TermTypeM Size
check e :: Size
e@(AppExp (Match Size
_ NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
_) = do
      let ps :: NonEmpty (PatBase Info VName StructType)
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 Size
_ SrcLoc
_) -> PatBase Info VName StructType
p) NonEmpty (CaseBase Info VName)
cs
      case [PatBase Info VName StructType] -> [Match ()]
unmatched ([PatBase Info VName StructType] -> [Match ()])
-> [PatBase Info VName StructType] -> [Match ()]
forall a b. (a -> b) -> a -> b
$ NonEmpty (PatBase Info VName StructType)
-> [PatBase Info VName StructType]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (PatBase Info VName StructType)
ps of
        [] -> Size -> TermTypeM Size
recurse Size
e
        [Match ()]
ps' ->
          SrcLoc -> Notes -> Doc () -> TermTypeM Size
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 Size)
-> (Doc () -> Doc ()) -> Doc () -> TermTypeM Size
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 Size) -> Doc () -> TermTypeM Size
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Unmatched cases in match expression:"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
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 :: Size
e@(AppExp (LetPat [SizeBinder VName]
_ PatBase Info VName StructType
p Size
_ Size
_ 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 Size -> TermTypeM Size
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Size -> TermTypeM Size
recurse Size
e
    check e :: Size
e@(Lambda [Pat ParamType]
ps Size
_ Maybe (TypeExp Size 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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) [Pat ParamType]
ps TermTypeM () -> TermTypeM Size -> TermTypeM Size
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Size -> TermTypeM Size
recurse Size
e
    check e :: Size
e@(AppExp (LetFun VName
_ ([TypeParamBase VName]
_, [Pat ParamType]
ps, Maybe (TypeExp Size VName)
_, Info ResRetType
_, Size
_) Size
_ SrcLoc
_) Info AppRes
_) =
      (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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) [Pat ParamType]
ps TermTypeM () -> TermTypeM Size -> TermTypeM Size
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Size -> TermTypeM Size
recurse Size
e
    check e :: Size
e@(AppExp (Loop [VName]
_ Pat ParamType
p LoopInitBase Info VName
_ LoopFormBase Info VName
form Size
_ 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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct Pat ParamType
p)
      case LoopFormBase Info VName
form of
        ForIn PatBase Info VName StructType
form_p Size
_ -> 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 ()
      Size -> TermTypeM Size
recurse Size
e
    check e :: Size
e@(IntLit Integer
x Info StructType
ty SrcLoc
loc) =
      Size
e Size -> TermTypeM () -> TermTypeM Size
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 :: Size
e@(FloatLit Double
x Info StructType
ty SrcLoc
loc) =
      Size
e Size -> TermTypeM () -> TermTypeM Size
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 :: Size
e@(Negate (IntLit Integer
x Info StructType
ty SrcLoc
loc1) SrcLoc
loc2) =
      Size
e Size -> TermTypeM () -> TermTypeM Size
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 e :: Size
e@(AppExp (BinOp (QualName [] VName
v, SrcLoc
_) Info StructType
_ (Size
x, Info (Maybe VName)
_) (Size, Info (Maybe VName))
_ SrcLoc
loc) Info AppRes
_)
      | VName -> Name
baseName VName
v Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"==",
        Array {} <- Size -> StructType
typeOf Size
x,
        VName -> Int
baseTag VName
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag = 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."
          Size -> TermTypeM Size
recurse Size
e
    check Size
e = Size -> TermTypeM Size
recurse Size
e
    recurse :: Size -> TermTypeM Size
recurse = ASTMapper TermTypeM -> Size -> TermTypeM Size
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Size -> m Size
astMap ASTMapper TermTypeM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = check}

    bitWidth :: IntType -> Int
bitWidth IntType
ty = Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* IntType -> Int
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 -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (IntType -> Int
bitWidth IntType
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2 a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (IntType -> Int
bitWidth IntType
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ IntType -> Int
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 ()
"."

-- | 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 Size VName), ResRetType, Size)
checkFunDef (VName
fname, Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl, [TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, ExpBase NoInfo VName
body, SrcLoc
loc) =
  (ExpBase NoInfo VName -> TermTypeM Size)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a.
(ExpBase NoInfo VName -> TermTypeM Size) -> TermTypeM a -> TypeM a
runTermTypeM ExpBase NoInfo VName -> TermTypeM Size
checkExp (TermTypeM
   ([TypeParamBase VName], [Pat ParamType],
    Maybe (TypeExp Size VName), ResRetType, Size)
 -> TypeM
      ([TypeParamBase VName], [Pat ParamType],
       Maybe (TypeExp Size VName), ResRetType, Size))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
-> TypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a b. (a -> b) -> a -> b
$ do
    ([TypeParamBase VName]
tparams', [Pat ParamType]
params', Maybe (TypeExp Size VName)
maybe_retdecl', RetType [VName]
dims TypeBase Size Uniqueness
rettype', Size
body') <-
      (VName, Maybe (TypeExp (ExpBase NoInfo VName) VName),
 [TypeParamBase VName], [PatBase NoInfo VName ParamType],
 ExpBase NoInfo VName, SrcLoc)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
checkBinding (VName
fname, Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl, [TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, ExpBase NoInfo VName
body, SrcLoc
loc)

    -- Since this is a top-level function, we also resolve overloaded
    -- types, using either defaults or complaining about ambiguities.
    Set VName -> TermTypeM ()
fixOverloadedTypes (Set VName -> TermTypeM ()) -> Set VName -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
      TypeBase Size Uniqueness -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars TypeBase Size Uniqueness
rettype' Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (Pat ParamType -> Set VName) -> [Pat ParamType] -> 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 (ParamType -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars (ParamType -> Set VName)
-> (Pat ParamType -> ParamType) -> Pat ParamType -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> ParamType
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType) [Pat ParamType]
params'

    -- Then replace all inferred types in the body and parameters.
    Size
body'' <- Size -> TermTypeM Size
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully Size
body'
    [Pat ParamType]
params'' <- (Pat ParamType -> TermTypeM (Pat ParamType))
-> [Pat ParamType] -> TermTypeM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> TermTypeM (Pat ParamType)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully [Pat ParamType]
params'
    Maybe (TypeExp Size VName)
maybe_retdecl'' <- (TypeExp Size VName -> TermTypeM (TypeExp Size VName))
-> Maybe (TypeExp Size VName)
-> TermTypeM (Maybe (TypeExp Size 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 Size VName -> TermTypeM (TypeExp Size VName)
forall e. ASTMappable e => e -> TermTypeM e
updateTypes Maybe (TypeExp Size VName)
maybe_retdecl'
    TypeBase Size Uniqueness
rettype'' <- TypeBase Size Uniqueness -> TermTypeM (TypeBase Size Uniqueness)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully TypeBase Size Uniqueness
rettype'

    -- Check if the function body can actually be evaluated.
    Size -> TermTypeM ()
causalityCheck Size
body''

    -- Check for various problems.
    (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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) [Pat ParamType]
params'
    Size -> TermTypeM ()
localChecks Size
body''

    let ((Size
body''', ResRetType
updated_ret), [TypeError]
errors) =
          (VName, [Pat ParamType], Size, ResRetType,
 Maybe (TypeExp Size VName), SrcLoc)
-> ((Size, ResRetType), [TypeError])
Consumption.checkValDef
            ( VName
fname,
              [Pat ParamType]
params'',
              Size
body'',
              [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Size Uniqueness
rettype'',
              Maybe (TypeExp Size VName)
maybe_retdecl'',
              SrcLoc
loc
            )

    (TypeError -> TermTypeM Any) -> [TypeError] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ TypeError -> TermTypeM Any
forall a. TypeError -> TermTypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError [TypeError]
errors

    ([TypeParamBase VName], [Pat ParamType],
 Maybe (TypeExp Size VName), ResRetType, Size)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TypeParamBase VName]
tparams', [Pat ParamType]
params'', Maybe (TypeExp Size VName)
maybe_retdecl'', ResRetType
updated_ret, Size
body''')

-- | This is "fixing" as in "setting them", not "correcting them".  We
-- only make very conservative fixing.
fixOverloadedTypes :: Names -> TermTypeM ()
fixOverloadedTypes :: Set VName -> TermTypeM ()
fixOverloadedTypes Set VName
tyvars_at_toplevel =
  TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints TermTypeM Constraints
-> (Constraints -> TermTypeM ()) -> TermTypeM ()
forall a b. TermTypeM a -> (a -> TermTypeM b) -> TermTypeM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ((VName, Constraint) -> TermTypeM ())
-> [(VName, Constraint)] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (VName, Constraint) -> TermTypeM ()
forall {m :: * -> *}.
(MonadUnify m, MonadTypeChecker m) =>
(VName, Constraint) -> m ()
fixOverloaded ([(VName, Constraint)] -> TermTypeM ())
-> (Constraints -> [(VName, Constraint)])
-> Constraints
-> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Constraint -> [(VName, Constraint)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName Constraint -> [(VName, Constraint)])
-> (Constraints -> Map VName Constraint)
-> Constraints
-> [(VName, Constraint)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Constraint) -> Constraint)
-> Constraints -> Map VName Constraint
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Int, Constraint) -> Constraint
forall a b. (a, b) -> b
snd
  where
    fixOverloaded :: (VName, Constraint) -> m ()
fixOverloaded (VName
v, Overloaded [PrimType]
ots Usage
usage)
      | IntType -> PrimType
Signed IntType
Int32 PrimType -> [PrimType] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PrimType]
ots = do
          Usage -> StructType -> StructType -> m ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify Usage
usage (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness
-> QualName VName
-> [TypeArg Size]
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar NoUniqueness
forall a. Monoid a => a
mempty (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) [])) (StructType -> m ()) -> StructType -> m ()
forall a b. (a -> b) -> a -> b
$
            ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Size NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Size NoUniqueness)
-> PrimType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int32)
          Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
tyvars_at_toplevel) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
            Usage -> Doc () -> m ()
forall loc. Located loc => loc -> Doc () -> m ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn Usage
usage Doc ()
"Defaulting ambiguous type to i32."
      | FloatType -> PrimType
FloatType FloatType
Float64 PrimType -> [PrimType] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PrimType]
ots = do
          Usage -> StructType -> StructType -> m ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify Usage
usage (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness
-> QualName VName
-> [TypeArg Size]
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar NoUniqueness
forall a. Monoid a => a
mempty (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) [])) (StructType -> m ()) -> StructType -> m ()
forall a b. (a -> b) -> a -> b
$
            ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Size NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Size NoUniqueness)
-> PrimType -> ScalarTypeBase Size NoUniqueness
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
FloatType FloatType
Float64)
          Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
tyvars_at_toplevel) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
            Usage -> Doc () -> m ()
forall loc. Located loc => loc -> Doc () -> m ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn Usage
usage Doc ()
"Defaulting ambiguous type to f64."
      | Bool
otherwise =
          Usage -> 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 Usage
usage 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-type" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Type is ambiguous (could be one of"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
commasep ((PrimType -> Doc ()) -> [PrimType] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. PrimType -> Doc ann
pretty [PrimType]
ots)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
")."
                Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Add a type annotation to disambiguate the type."
    fixOverloaded (VName
v, NoConstraint Liftedness
_ Usage
usage) = do
      -- See #1552.
      Usage -> StructType -> StructType -> m ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify Usage
usage (ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness
-> QualName VName
-> [TypeArg Size]
-> ScalarTypeBase Size NoUniqueness
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar NoUniqueness
forall a. Monoid a => a
mempty (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) [])) (StructType -> m ()) -> StructType -> m ()
forall a b. (a -> b) -> a -> b
$
        ScalarTypeBase Size NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ([StructType] -> ScalarTypeBase Size NoUniqueness
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord [])
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
tyvars_at_toplevel) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
        Usage -> Doc () -> m ()
forall loc. Located loc => loc -> Doc () -> m ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn Usage
usage Doc ()
"Defaulting ambiguous type to ()."
    fixOverloaded (VName
_, Equality Usage
usage) =
      Usage -> 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 Usage
usage 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-type" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Type is ambiguous (must be equality type)."
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Add a type annotation to disambiguate the type."
    fixOverloaded (VName
_, HasFields Liftedness
_ Map Name StructType
fs Usage
usage) =
      Usage -> 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 Usage
usage 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-type" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Type is ambiguous.  Must be record with fields:"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
stack ([Doc ()] -> Doc ()) -> [Doc ()] -> Doc ()
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> Doc ()) -> [(Name, StructType)] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> Doc ()
forall {a} {a} {ann}. (Pretty a, Pretty a) => (a, a) -> Doc ann
field ([(Name, StructType)] -> [Doc ()])
-> [(Name, StructType)] -> [Doc ()]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name StructType
fs)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Add a type annotation to disambiguate the type."
      where
        field :: (a, a) -> Doc ann
field (a
l, a
t) = a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
l Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align (a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
t)
    fixOverloaded (VName
_, HasConstrs Liftedness
_ Map Name [StructType]
cs Usage
usage) =
      Usage -> 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 Usage
usage 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-type" (Doc () -> m ()) -> Doc () -> m ()
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Type is ambiguous (must be a sum type with constructors:"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> ScalarTypeBase Size NoUniqueness -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ScalarTypeBase Size NoUniqueness -> Doc ann
pretty (Map Name [StructType] -> ScalarTypeBase Size NoUniqueness
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum Map Name [StructType]
cs)
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
")."
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Add a type annotation to disambiguate the type."
    fixOverloaded (VName
v, Size Maybe Size
Nothing (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 ()
"."
    fixOverloaded (VName
v, Size Maybe Size
Nothing (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 ()
"."
    fixOverloaded (VName, Constraint)
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

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)
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

inferredReturnType :: SrcLoc -> [Pat ParamType] -> StructType -> TermTypeM StructType
inferredReturnType :: SrcLoc -> [Pat ParamType] -> StructType -> TermTypeM StructType
inferredReturnType SrcLoc
loc [Pat ParamType]
params StructType
t = do
  -- The inferred type may refer to names that are bound by the
  -- parameter patterns, but which will not be visible in the type.
  -- These we must turn into fresh type variables, which will be
  -- existential in the return type.
  (StructType, [VName]) -> StructType
forall a b. (a, b) -> a
fst ((StructType, [VName]) -> StructType)
-> TermTypeM (StructType, [VName]) -> TermTypeM StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc -> [VName] -> StructType -> TermTypeM (StructType, [VName])
forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType SrcLoc
loc [VName]
hidden_params StructType
t
  where
    hidden_params :: [VName]
hidden_params = (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]
hidden) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params
    hidden :: [VName]
hidden = [Pat ParamType] -> [VName]
hiddenParamNames [Pat ParamType]
params

checkBinding ::
  ( VName,
    Maybe (TypeExp (ExpBase NoInfo VName) VName),
    [TypeParam],
    [PatBase NoInfo VName ParamType],
    ExpBase NoInfo VName,
    SrcLoc
  ) ->
  TermTypeM
    ( [TypeParam],
      [Pat ParamType],
      Maybe (TypeExp Exp VName),
      ResRetType,
      Exp
    )
checkBinding :: (VName, Maybe (TypeExp (ExpBase NoInfo VName) VName),
 [TypeParamBase VName], [PatBase NoInfo VName ParamType],
 ExpBase NoInfo VName, SrcLoc)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
checkBinding (VName
fname, Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl, [TypeParamBase VName]
tparams, [PatBase NoInfo VName ParamType]
params, ExpBase NoInfo VName
body, SrcLoc
loc) =
  TermTypeM
  ([TypeParamBase VName], [Pat ParamType],
   Maybe (TypeExp Size VName), ResRetType, Size)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a. TermTypeM a -> TermTypeM a
incLevel (TermTypeM
   ([TypeParamBase VName], [Pat ParamType],
    Maybe (TypeExp Size VName), ResRetType, Size)
 -> TermTypeM
      ([TypeParamBase VName], [Pat ParamType],
       Maybe (TypeExp Size VName), ResRetType, Size))
-> (([Pat ParamType]
     -> TermTypeM
          ([TypeParamBase VName], [Pat ParamType],
           Maybe (TypeExp Size VName), ResRetType, Size))
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType],
          Maybe (TypeExp Size VName), ResRetType, Size))
-> ([Pat ParamType]
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType],
          Maybe (TypeExp Size VName), ResRetType, Size))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeParamBase VName]
-> [PatBase NoInfo VName ParamType]
-> ([Pat ParamType]
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType],
          Maybe (TypeExp Size VName), ResRetType, Size))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a.
[TypeParamBase VName]
-> [PatBase NoInfo VName ParamType]
-> ([Pat ParamType] -> TermTypeM a)
-> TermTypeM a
bindingParams [TypeParamBase VName]
tparams [PatBase NoInfo VName ParamType]
params (([Pat ParamType]
  -> TermTypeM
       ([TypeParamBase VName], [Pat ParamType],
        Maybe (TypeExp Size VName), ResRetType, Size))
 -> TermTypeM
      ([TypeParamBase VName], [Pat ParamType],
       Maybe (TypeExp Size VName), ResRetType, Size))
-> ([Pat ParamType]
    -> TermTypeM
         ([TypeParamBase VName], [Pat ParamType],
          Maybe (TypeExp Size VName), ResRetType, Size))
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
    Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
maybe_retdecl' <- (TypeExp (ExpBase NoInfo VName) VName
 -> TermTypeM
      (TypeExp Size VName, TypeBase Size Uniqueness, [VName]))
-> Maybe (TypeExp (ExpBase NoInfo VName) VName)
-> TermTypeM
     (Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [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 (ExpBase NoInfo VName) VName
-> TermTypeM
     (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
checkTypeExpNonrigid Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl

    Size
body' <-
      [Pat ParamType]
-> ExpBase NoInfo VName
-> Maybe (TypeBase Size Uniqueness)
-> SrcLoc
-> TermTypeM Size
checkFunBody
        [Pat ParamType]
params'
        ExpBase NoInfo VName
body
        ((\(TypeExp Size VName
_, TypeBase Size Uniqueness
x, [VName]
_) -> TypeBase Size Uniqueness
x) ((TypeExp Size VName, TypeBase Size Uniqueness, [VName])
 -> TypeBase Size Uniqueness)
-> Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
-> Maybe (TypeBase Size Uniqueness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
maybe_retdecl')
        (SrcLoc
-> (TypeExp (ExpBase NoInfo VName) VName -> SrcLoc)
-> Maybe (TypeExp (ExpBase NoInfo VName) VName)
-> SrcLoc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe SrcLoc
loc TypeExp (ExpBase NoInfo VName) VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Maybe (TypeExp (ExpBase NoInfo VName) VName)
maybe_retdecl)

    [Pat ParamType]
params'' <- (Pat ParamType -> TermTypeM (Pat ParamType))
-> [Pat ParamType] -> TermTypeM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> TermTypeM (Pat ParamType)
forall e. ASTMappable e => e -> TermTypeM e
updateTypes [Pat ParamType]
params'
    StructType
body_t <- Size -> TermTypeM StructType
expTypeFully Size
body'

    (Maybe (TypeExp Size VName)
maybe_retdecl'', TypeBase Size Uniqueness
rettype) <- case Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
maybe_retdecl' of
      Just (TypeExp Size VName
retdecl', TypeBase Size Uniqueness
ret, [VName]
_) -> do
        TypeBase Size Uniqueness
ret' <- TypeBase Size Uniqueness -> TermTypeM (TypeBase Size Uniqueness)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully TypeBase Size Uniqueness
ret
        (Maybe (TypeExp Size VName), TypeBase Size Uniqueness)
-> TermTypeM (Maybe (TypeExp Size VName), TypeBase Size Uniqueness)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp Size VName -> Maybe (TypeExp Size VName)
forall a. a -> Maybe a
Just TypeExp Size VName
retdecl', TypeBase Size Uniqueness
ret')
      Maybe (TypeExp Size VName, TypeBase Size Uniqueness, [VName])
Nothing
        | [PatBase NoInfo VName ParamType] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatBase NoInfo VName ParamType]
params ->
            (Maybe (TypeExp Size VName), TypeBase Size Uniqueness)
-> TermTypeM (Maybe (TypeExp Size VName), TypeBase Size Uniqueness)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (TypeExp Size VName)
forall a. Maybe a
Nothing, Uniqueness -> StructType -> TypeBase Size Uniqueness
forall u. Uniqueness -> TypeBase Size u -> TypeBase Size Uniqueness
toRes Uniqueness
Nonunique StructType
body_t)
        | Bool
otherwise -> do
            StructType
body_t' <- SrcLoc -> [Pat ParamType] -> StructType -> TermTypeM StructType
inferredReturnType SrcLoc
loc [Pat ParamType]
params'' StructType
body_t
            (Maybe (TypeExp Size VName), TypeBase Size Uniqueness)
-> TermTypeM (Maybe (TypeExp Size VName), TypeBase Size Uniqueness)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (TypeExp Size VName)
forall a. Maybe a
Nothing, Uniqueness -> StructType -> TypeBase Size Uniqueness
forall u. Uniqueness -> TypeBase Size u -> TypeBase Size Uniqueness
toRes Uniqueness
Nonunique StructType
body_t')

    Maybe VName -> [Pat ParamType] -> TermTypeM ()
verifyFunctionParams (VName -> Maybe VName
forall a. a -> Maybe a
Just VName
fname) [Pat ParamType]
params''

    ([TypeParamBase VName]
tparams', [Pat ParamType]
params''', ResRetType
rettype') <-
      Name
-> SrcLoc
-> [TypeParamBase VName]
-> [Pat ParamType]
-> TypeBase Size Uniqueness
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
letGeneralise (VName -> Name
baseName VName
fname) SrcLoc
loc [TypeParamBase VName]
tparams [Pat ParamType]
params'' (TypeBase Size Uniqueness
 -> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType))
-> TermTypeM (TypeBase Size Uniqueness)
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TypeBase Size Uniqueness -> TermTypeM (TypeBase Size Uniqueness)
forall u. TypeBase Size u -> TermTypeM (TypeBase Size u)
unscopeUnknown TypeBase Size Uniqueness
rettype

    Bool -> TermTypeM () -> TermTypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
      ( [PatBase NoInfo VName ParamType] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatBase NoInfo VName ParamType]
params
          Bool -> Bool -> Bool
&& (TypeParamBase VName -> Bool) -> [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isSizeParam [TypeParamBase VName]
tparams'
          Bool -> Bool -> Bool
&& Bool -> Bool
not ([VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (ResRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims ResRetType
rettype'))
      )
      (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ 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
$ Text -> Doc ()
forall a. Text -> Doc a
textwrap Text
"A size-polymorphic value binding may not have a type with an existential size."
        Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Type of this binding is:"
        Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ResRetType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ResRetType -> Doc ann
pretty ResRetType
rettype')
        Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"with the following type parameters:"
        Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
sep ([Doc ()] -> Doc ()) -> [Doc ()] -> Doc ()
forall a b. (a -> b) -> a -> b
$ (TypeParamBase VName -> Doc ())
-> [TypeParamBase VName] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeParamBase VName -> Doc ann
pretty ([TypeParamBase VName] -> [Doc ()])
-> [TypeParamBase VName] -> [Doc ()]
forall a b. (a -> b) -> a -> b
$ (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isSizeParam [TypeParamBase VName]
tparams')

    ([TypeParamBase VName], [Pat ParamType],
 Maybe (TypeExp Size VName), ResRetType, Size)
-> TermTypeM
     ([TypeParamBase VName], [Pat ParamType],
      Maybe (TypeExp Size VName), ResRetType, Size)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TypeParamBase VName]
tparams', [Pat ParamType]
params''', Maybe (TypeExp Size VName)
maybe_retdecl'', ResRetType
rettype', Size
body')

-- | Extract all the shape names that occur in positive position
-- (roughly, left side of an arrow) in a given type.
sizeNamesPos :: TypeBase Size als -> S.Set VName
sizeNamesPos :: forall als. TypeBase Size als -> Set VName
sizeNamesPos (Scalar (Arrow als
_ PName
_ Diet
_ StructType
t1 (RetType [VName]
_ TypeBase Size Uniqueness
t2))) = StructType -> Set VName
forall als. TypeBase Size als -> Set VName
onParam StructType
t1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> TypeBase Size Uniqueness -> Set VName
forall als. TypeBase Size als -> Set VName
sizeNamesPos TypeBase Size Uniqueness
t2
  where
    onParam :: TypeBase Size als -> S.Set VName
    onParam :: forall als. TypeBase Size als -> Set VName
onParam (Scalar Arrow {}) = Set VName
forall a. Monoid a => a
mempty
    onParam (Scalar (Record Map Name (TypeBase Size als)
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 Size als -> Set VName)
-> [TypeBase Size als] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Size als -> Set VName
forall als. TypeBase Size als -> Set VName
onParam ([TypeBase Size als] -> [Set VName])
-> [TypeBase Size als] -> [Set VName]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Size als) -> [TypeBase Size als]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase Size als)
fs
    onParam (Scalar (TypeVar als
_ QualName VName
_ [TypeArg Size]
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 Size -> Set VName) -> [TypeArg Size] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg Size -> Set VName
onTypeArg [TypeArg Size]
targs
    onParam TypeBase Size als
t = FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ TypeBase Size als -> FV
forall u. TypeBase Size u -> FV
freeInType TypeBase Size als
t
    onTypeArg :: TypeArg Size -> 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 Size
_) = Set VName
forall a. Monoid a => a
mempty
    onTypeArg (TypeArgType StructType
t) = StructType -> Set VName
forall als. TypeBase Size als -> Set VName
onParam StructType
t
sizeNamesPos TypeBase Size als
_ = 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.
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] -> TermTypeM ())
-> TermTypeM [Pat ParamType] -> TermTypeM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Pat ParamType -> TermTypeM (Pat ParamType))
-> [Pat ParamType] -> TermTypeM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> TermTypeM (Pat ParamType)
forall e. ASTMappable e => e -> TermTypeM e
updateTypes [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 u. Pat (TypeBase Size u) -> 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)
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 u -> RetTypeBase Size u
injectExt :: forall u. [VName] -> TypeBase Size u -> RetTypeBase Size u
injectExt [] TypeBase Size u
ret = [VName] -> TypeBase Size u -> RetTypeBase Size u
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Size u
ret
injectExt [VName]
ext TypeBase Size u
ret = [VName] -> TypeBase Size u -> RetTypeBase Size u
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext_here (TypeBase Size u -> RetTypeBase Size u)
-> TypeBase Size u -> RetTypeBase Size u
forall a b. (a -> b) -> a -> b
$ TypeBase Size u -> TypeBase Size u
forall u. TypeBase Size u -> TypeBase Size u
deeper TypeBase Size u
ret
  where
    (Set VName
immediate, Set VName
_) = TypeBase Size u -> (Set VName, Set VName)
forall u. TypeBase Size u -> (Set VName, Set VName)
dimUses TypeBase Size u
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 u -> TypeBase Size u
    deeper :: forall u. TypeBase Size u -> TypeBase Size u
deeper (Scalar (Prim PrimType
t)) = ScalarTypeBase Size u -> TypeBase Size u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size u -> TypeBase Size u)
-> ScalarTypeBase Size u -> TypeBase Size u
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t
    deeper (Scalar (Record Map Name (TypeBase Size u)
fs)) = ScalarTypeBase Size u -> TypeBase Size u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size u -> TypeBase Size u)
-> ScalarTypeBase Size u -> TypeBase Size u
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Size u) -> ScalarTypeBase Size u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase Size u) -> ScalarTypeBase Size u)
-> Map Name (TypeBase Size u) -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ (TypeBase Size u -> TypeBase Size u)
-> Map Name (TypeBase Size u) -> Map Name (TypeBase Size u)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase Size u -> TypeBase Size u
forall u. TypeBase Size u -> TypeBase Size u
deeper Map Name (TypeBase Size u)
fs
    deeper (Scalar (Sum Map Name [TypeBase Size u]
cs)) = ScalarTypeBase Size u -> TypeBase Size u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size u -> TypeBase Size u)
-> ScalarTypeBase Size u -> TypeBase Size u
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase Size u] -> ScalarTypeBase Size u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase Size u] -> ScalarTypeBase Size u)
-> Map Name [TypeBase Size u] -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ ([TypeBase Size u] -> [TypeBase Size u])
-> Map Name [TypeBase Size u] -> Map Name [TypeBase Size u]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase Size u -> TypeBase Size u)
-> [TypeBase Size u] -> [TypeBase Size u]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Size u -> TypeBase Size u
forall u. TypeBase Size u -> TypeBase Size u
deeper) Map Name [TypeBase Size u]
cs
    deeper (Scalar (Arrow u
als PName
p Diet
d1 StructType
t1 (RetType [VName]
t2_ext TypeBase Size Uniqueness
t2))) =
      ScalarTypeBase Size u -> TypeBase Size u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size u -> TypeBase Size u)
-> ScalarTypeBase Size u -> TypeBase Size u
forall a b. (a -> b) -> a -> b
$ u
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Size u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
als PName
p Diet
d1 StructType
t1 (ResRetType -> ScalarTypeBase Size u)
-> ResRetType -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall u. [VName] -> TypeBase Size u -> RetTypeBase Size u
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)) TypeBase Size Uniqueness
t2
    deeper (Scalar (TypeVar u
u QualName VName
tn [TypeArg Size]
targs)) =
      ScalarTypeBase Size u -> TypeBase Size u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size u -> TypeBase Size u)
-> ScalarTypeBase Size u -> TypeBase Size u
forall a b. (a -> b) -> a -> b
$ u -> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size u
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar u
u QualName VName
tn ([TypeArg Size] -> ScalarTypeBase Size u)
-> [TypeArg Size] -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ (TypeArg Size -> TypeArg Size) -> [TypeArg Size] -> [TypeArg Size]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg Size -> TypeArg Size
deeperArg [TypeArg Size]
targs
    deeper t :: TypeBase Size u
t@Array {} = TypeBase Size u
t

    deeperArg :: TypeArg Size -> TypeArg Size
deeperArg (TypeArgType StructType
t) = StructType -> TypeArg Size
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (StructType -> TypeArg Size) -> StructType -> TypeArg Size
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall u. TypeBase Size u -> TypeBase Size u
deeper StructType
t
    deeperArg (TypeArgDim Size
d) = Size -> TypeArg Size
forall dim. dim -> TypeArg dim
TypeArgDim Size
d

-- | Find all type variables in the given type that are covered by the
-- constraints, and produce type parameters that close over them.
--
-- The passed-in list of type parameters is always prepended to the
-- produced list of type parameters.
closeOverTypes ::
  Name ->
  SrcLoc ->
  [TypeParam] ->
  [StructType] ->
  ResType ->
  Constraints ->
  TermTypeM ([TypeParam], ResRetType)
closeOverTypes :: Name
-> SrcLoc
-> [TypeParamBase VName]
-> [StructType]
-> TypeBase Size Uniqueness
-> Constraints
-> TermTypeM ([TypeParamBase VName], ResRetType)
closeOverTypes Name
defname SrcLoc
defloc [TypeParamBase VName]
tparams [StructType]
paramts TypeBase Size Uniqueness
ret Constraints
substs = do
  ([TypeParamBase VName]
more_tparams, [VName]
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
$ ((Int, Constraint) -> Constraint)
-> Constraints -> Map VName Constraint
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Int, Constraint) -> Constraint
forall a b. (a, b) -> b
snd Constraints
to_close_over)
  let mkExt :: VName -> Maybe VName
mkExt VName
v =
        case VName -> Constraints -> Maybe (Int, Constraint)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Constraints
substs of
          Just (Int
_, UnknownSize {}) -> VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
          Maybe (Int, Constraint)
_ -> Maybe VName
forall a. Maybe a
Nothing
  ([TypeParamBase VName], ResRetType)
-> TermTypeM ([TypeParamBase VName], ResRetType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( [TypeParamBase VName]
tparams [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
more_tparams,
      [VName] -> TypeBase Size Uniqueness -> ResRetType
forall u. [VName] -> TypeBase Size u -> RetTypeBase Size u
injectExt ([VName] -> [VName]
forall a. Ord a => [a] -> [a]
nubOrd ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ [VName]
retext [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ (VName -> Maybe VName) -> [VName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe VName
mkExt (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
$ TypeBase Size Uniqueness -> FV
forall u. TypeBase Size u -> FV
freeInType TypeBase Size Uniqueness
ret)) TypeBase Size Uniqueness
ret
    )
  where
    -- Diet does not matter here.
    t :: StructType
t = [ParamType] -> ResRetType -> StructType
foldFunType ((StructType -> ParamType) -> [StructType] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Size u -> ParamType
toParam Diet
Observe) [StructType]
paramts) (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Size Uniqueness
ret
    to_close_over :: Constraints
to_close_over = (VName -> (Int, Constraint) -> Bool) -> Constraints -> Constraints
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\VName
k (Int, Constraint)
_ -> VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
visible) Constraints
substs
    visible :: Set VName
visible = StructType -> Set VName
forall dim as. TypeBase dim as -> 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 u. TypeBase Size u -> FV
freeInType StructType
t)

    (Set VName
produced_sizes, Set VName
param_sizes) = StructType -> (Set VName, Set VName)
forall u. TypeBase Size u -> (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, NoConstraint Liftedness
l Usage
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
$ Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
l VName
k (SrcLoc -> TypeParamBase VName) -> SrcLoc -> TypeParamBase VName
forall a b. (a -> b) -> a -> b
$ Usage -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Usage
usage
    closeOver (VName
k, ParamType Liftedness
l Loc
loc) =
      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
$ Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
l VName
k (SrcLoc -> TypeParamBase VName) -> SrcLoc -> TypeParamBase VName
forall a b. (a -> b) -> a -> b
$ Loc -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Loc
loc
    closeOver (VName
k, Size Maybe Size
Nothing Usage
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 -> TypeParamBase VName) -> SrcLoc -> TypeParamBase VName
forall a b. (a -> b) -> a -> b
$ Usage -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Usage
usage
    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
notes <- SrcLoc -> Size -> f Notes
forall a (m :: * -> *).
(Located a, MonadUnify m) =>
a -> Size -> m Notes
dimNotes SrcLoc
defloc (Size -> f Notes) -> Size -> f Notes
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Size
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
k) SrcLoc
forall a. Monoid a => a
mempty
          SrcLoc
-> Notes
-> Doc ()
-> f (Maybe (Either (TypeParamBase VName) VName))
forall loc a. Located loc => loc -> Notes -> Doc () -> f a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
defloc Notes
notes (Doc () -> f (Maybe (Either (TypeParamBase VName) VName)))
-> (Doc () -> Doc ())
-> Doc ()
-> f (Maybe (Either (TypeParamBase VName) VName))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"unknown-param-def" (Doc () -> f (Maybe (Either (TypeParamBase VName) VName)))
-> Doc () -> f (Maybe (Either (TypeParamBase VName) VName))
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Unknown 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
k)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"in parameter of"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Name -> Doc ()
forall ann. Name -> Doc ann
forall v a. IsName v => v -> Doc a
prettyName Name
defname)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
", which is inferred as:"
                Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (StructType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty StructType
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]
-> TypeBase Size Uniqueness
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
letGeneralise Name
defname SrcLoc
defloc [TypeParamBase VName]
tparams [Pat ParamType]
params TypeBase Size Uniqueness
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
    Constraints
now_substs <- TermTypeM Constraints
forall (m :: * -> *). MonadUnify m => m Constraints
getConstraints

    -- Candidates for let-generalisation are those type variables that
    --
    -- (1) were not known before we checked this function, and
    --
    -- (2) are not used in the (new) definition of any type variables
    -- known before we checked this function.
    --
    -- (3) are not referenced from an overloaded type (for example,
    -- are the element types of an incompletely resolved record type).
    -- This is a bit more restrictive than I'd like, and SML for
    -- example does not have this restriction.
    --
    -- Criteria (1) and (2) is implemented by looking at the binding
    -- level of the type variables.
    let keep_type_vars :: Set VName
keep_type_vars = Constraints -> Set VName
overloadedTypeVars Constraints
now_substs

    Int
cur_lvl <- TermTypeM Int
forall (m :: * -> *). MonadUnify m => m Int
curLevel
    let candidate :: VName -> (Int, b) -> Bool
candidate VName
k (Int
lvl, b
_) = (VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
keep_type_vars) Bool -> Bool -> Bool
&& Int
lvl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= (Int
cur_lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Pat ParamType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Pat ParamType]
params)
        new_substs :: Constraints
new_substs = (VName -> (Int, Constraint) -> Bool) -> Constraints -> Constraints
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey VName -> (Int, Constraint) -> Bool
forall {b}. VName -> (Int, b) -> Bool
candidate Constraints
now_substs

    ([TypeParamBase VName]
tparams', RetType [VName]
ret_dims TypeBase Size Uniqueness
restype') <-
      Name
-> SrcLoc
-> [TypeParamBase VName]
-> [StructType]
-> TypeBase Size Uniqueness
-> Constraints
-> TermTypeM ([TypeParamBase VName], ResRetType)
closeOverTypes
        Name
defname
        SrcLoc
defloc
        [TypeParamBase VName]
tparams
        ((Pat ParamType -> StructType) -> [Pat ParamType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map Pat ParamType -> StructType
forall u. Pat (TypeBase Size u) -> StructType
patternStructType [Pat ParamType]
params)
        TypeBase Size Uniqueness
restype
        Constraints
new_substs

    TypeBase Size Uniqueness
restype'' <- TypeBase Size Uniqueness -> TermTypeM (TypeBase Size Uniqueness)
forall e. ASTMappable e => e -> TermTypeM e
updateTypes TypeBase Size Uniqueness
restype'

    let used_sizes :: FV
used_sizes =
          TypeBase Size Uniqueness -> FV
forall u. TypeBase Size u -> FV
freeInType TypeBase Size Uniqueness
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 u. TypeBase Size u -> 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 u. Pat (TypeBase d u) -> TypeBase d u
patternType) [Pat ParamType]
params
    case (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` FV -> Set VName
fvVars FV
used_sizes) (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] -> [TypeParamBase VName])
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> a -> b
$
      (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isSizeParam [TypeParamBase VName]
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.
    (Constraints -> Constraints) -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
(Constraints -> Constraints) -> m ()
modifyConstraints ((Constraints -> Constraints) -> TermTypeM ())
-> (Constraints -> Constraints) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ (VName -> (Int, Constraint) -> Bool) -> Constraints -> Constraints
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey ((VName -> (Int, Constraint) -> Bool)
 -> Constraints -> Constraints)
-> (VName -> (Int, Constraint) -> Bool)
-> Constraints
-> Constraints
forall a b. (a -> b) -> a -> b
$ \VName
k (Int, Constraint)
_ -> VName
k VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` (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'

    ([TypeParamBase VName], [Pat ParamType], ResRetType)
-> TermTypeM ([TypeParamBase VName], [Pat ParamType], ResRetType)
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TypeParamBase VName]
tparams', [Pat ParamType]
params, [VName] -> TypeBase Size Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims TypeBase Size Uniqueness
restype'')

checkFunBody ::
  [Pat ParamType] ->
  ExpBase NoInfo VName ->
  Maybe ResType ->
  SrcLoc ->
  TermTypeM Exp
checkFunBody :: [Pat ParamType]
-> ExpBase NoInfo VName
-> Maybe (TypeBase Size Uniqueness)
-> SrcLoc
-> TermTypeM Size
checkFunBody [Pat ParamType]
params ExpBase NoInfo VName
body Maybe (TypeBase Size Uniqueness)
maybe_rettype SrcLoc
loc = do
  Size
body' <- ExpBase NoInfo VName -> TermTypeM Size
checkExp ExpBase NoInfo VName
body

  -- Unify body return type with return annotation, if one exists.
  case Maybe (TypeBase Size Uniqueness)
maybe_rettype of
    Just TypeBase Size Uniqueness
rettype -> do
      StructType
body_t <- Size -> TermTypeM StructType
expTypeFully Size
body'
      -- We need to turn any sizes provided by "hidden" parameter
      -- names into existential sizes instead.
      let hidden :: [VName]
hidden = [Pat ParamType] -> [VName]
hiddenParamNames [Pat ParamType]
params
      (StructType
body_t', [VName]
_) <-
        SrcLoc -> [VName] -> StructType -> TermTypeM (StructType, [VName])
forall as.
SrcLoc
-> [VName]
-> TypeBase Size as
-> TermTypeM (TypeBase Size as, [VName])
unscopeType
          SrcLoc
loc
          ((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]
hidden) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)
          StructType
body_t

      let usage :: Usage
usage = ExpBase NoInfo VName -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage ExpBase NoInfo VName
body Text
"return type annotation"
      Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure (TypeBase Size Uniqueness -> StructType -> Checking
CheckingReturn TypeBase Size Uniqueness
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 (TypeBase Size Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Size Uniqueness
rettype) StructType
body_t'
    Maybe (TypeBase Size Uniqueness)
Nothing -> () -> TermTypeM ()
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  Size -> TermTypeM Size
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
body'

arrayOfM ::
  SrcLoc ->
  StructType ->
  Shape Size ->
  TermTypeM StructType
arrayOfM :: SrcLoc -> StructType -> Shape Size -> TermTypeM StructType
arrayOfM SrcLoc
loc StructType
t Shape Size
shape = do
  Usage -> Text -> StructType -> TermTypeM ()
forall (m :: * -> *) dim u.
(MonadUnify m, Pretty (Shape dim), Pretty u) =>
Usage -> Text -> TypeBase dim u -> m ()
arrayElemType (SrcLoc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc Text
"use as array element") Text
"type used in array" StructType
t
  StructType -> TermTypeM StructType
forall a. a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType -> TermTypeM StructType)
-> StructType -> TermTypeM StructType
forall a b. (a -> b) -> a -> b
$ Shape Size -> StructType -> StructType
forall dim.
Shape dim -> TypeBase dim NoUniqueness -> TypeBase dim NoUniqueness
arrayOf Shape Size
shape StructType
t