content: -- | 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 Data.Text qualified as T
  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 Lambda {} = True
  hasBinding (AppExp LetPat {} _) = True
  hasBinding (AppExp LetFun {} _) = True
  hasBinding (AppExp Loop {} _) = True
  hasBinding (AppExp LetWith {} _) = True
  hasBinding (AppExp Match {} _) = True
  hasBinding e = isNothing $ astMap m e
    where
      m =
        identityMapper {mapOnExp = \e' -> if hasBinding e' then Nothing else Just e'}
  
  overloadedTypeVars :: Constraints -> Names
  overloadedTypeVars = mconcat . map f . M.elems
    where
      f (_, HasFields _ fs _) = mconcat $ map typeVars $ M.elems fs
      f _ = 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 loc t1 t2 =
    onFailure (CheckingBranches t1 t2) $
      unifyMostCommon (mkUsage loc "unification of branch results") t1 t2
  
  unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (StructType, [VName])
  unifyBranches loc e1 e2 = do
    e1_t <- expTypeFully e1
    e2_t <- expTypeFully e2
    unifyBranchTypes loc e1_t e2_t
  
  sliceShape ::
    Maybe (SrcLoc, Rigidity) ->
    [DimIndex] ->
    TypeBase Size as ->
    TermTypeM (TypeBase Size as, [VName])
  sliceShape r slice t@(Array u (Shape orig_dims) et) =
    runStateT (setDims <$> adjustDims slice orig_dims) []
    where
      setDims [] = stripArray (length orig_dims) t
      setDims dims' = Array u (Shape dims') 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 Rigid {} = True
      isRigid _ = False
      refine_sizes = maybe False (isRigid . snd) r
  
      sliceSize orig_d i j stride =
        case r of
          Just (loc, Rigid _) -> do
            (d, ext) <-
              lift . extSize loc $
                SourceSlice orig_d' (bareExp <$> i) (bareExp <$> j) (bareExp <$> stride)
            modify (maybeToList ext ++)
            pure d
          Just (loc, Nonrigid) ->
            lift $
              flip sizeFromName loc . qualName
                <$> newFlexibleDim (mkUsage loc "size of slice") "slice_dim"
          Nothing -> do
            v <- lift $ newID "slice_anydim"
            modify (v :)
            pure $ sizeFromName (qualName v) mempty
        where
          -- The original size does not matter if the slice is fully specified.
          orig_d'
            | isJust i, isJust j = Nothing
            | otherwise = Just orig_d
  
      warnIfBinding binds d i j stride size =
        if binds
          then do
            lift . warn (srclocOf size) $
              withIndexLink
                "size-expression-bind"
                "Size expression with binding is replaced by unknown size."
            (:) <$> sliceSize d i j stride
          else pure (size :)
  
      adjustDims (DimFix {} : idxes') (_ : dims) =
        adjustDims idxes' dims
      -- Pat match some known slices to be non-existential.
      adjustDims (DimSlice i j stride : idxes') (d : dims)
        | refine_sizes,
          maybe True ((== Just 0) . isInt64) i,
          maybe True ((== Just 1) . isInt64) stride = do
            let binds = maybe False hasBinding j
            warnIfBinding binds d i j stride (fromMaybe d j)
              <*> adjustDims idxes' dims
      adjustDims ((DimSlice i j stride) : idxes') (d : dims)
        | refine_sizes,
          Just i' <- i, -- if i ~ 0, previous case
          maybe True ((== Just 1) . isInt64) stride = do
            let j' = fromMaybe d j
                binds = hasBinding j' || hasBinding i'
            warnIfBinding binds d i j stride (sizeMinus j' i')
              <*> adjustDims idxes' dims
      -- stride == -1
      adjustDims ((DimSlice Nothing Nothing stride) : idxes') (d : dims)
        | refine_sizes,
          maybe True ((== Just (-1)) . isInt64) stride =
            (d :) <$> adjustDims idxes' dims
      adjustDims ((DimSlice (Just i) (Just j) stride) : idxes') (d : dims)
        | refine_sizes,
          maybe True ((== Just (-1)) . isInt64) stride = do
            let binds = hasBinding i || hasBinding j
            warnIfBinding binds d (Just i) (Just j) stride (sizeMinus i j)
              <*> adjustDims idxes' dims
      -- existential
      adjustDims ((DimSlice i j stride) : idxes') (d : dims) =
        (:) <$> sliceSize d i j stride <*> adjustDims idxes' dims
      adjustDims _ dims =
        pure dims
  
      sizeMinus j i =
        AppExp
          ( BinOp
              (qualName (intrinsicVar "-"), mempty)
              sizeBinOpInfo
              (j, Info Nothing)
              (i, Info Nothing)
              mempty
          )
          $ Info
          $ AppRes i64 []
      i64 = Scalar $ Prim $ Signed Int64
      sizeBinOpInfo = Info $ foldFunType [i64, i64] $ RetType [] i64
  sliceShape _ _ t = pure (t, [])
  
  --- Main checkers
  
  checkAscript ::
    SrcLoc ->
    TypeExp (ExpBase NoInfo VName) VName ->
    ExpBase NoInfo VName ->
    TermTypeM (TypeExp Exp VName, Exp)
  checkAscript loc te e = do
    (te', decl_t, _) <- checkTypeExpNonrigid te
    e' <- checkExp e
    e_t <- expTypeFully e'
  
    onFailure (CheckingAscription (toStruct decl_t) e_t) $
      unify (mkUsage loc "type ascription") (toStruct decl_t) e_t
  
    pure (te', e')
  
  checkCoerce ::
    SrcLoc ->
    TypeExp (ExpBase NoInfo VName) VName ->
    ExpBase NoInfo VName ->
    TermTypeM (TypeExp Exp VName, StructType, Exp)
  checkCoerce loc te e = do
    (te', te_t, ext) <- checkTypeExpNonrigid te
    e' <- checkExp e
    e_t <- expTypeFully e'
  
    te_t_nonrigid <- makeNonExtFresh ext $ toStruct te_t
  
    onFailure (CheckingAscription (toStruct te_t) e_t) $
      unify (mkUsage loc "size coercion") e_t te_t_nonrigid
  
    -- If the type expression had any anonymous dimensions, these will
    -- now be in 'ext'.  Those we keep nonrigid and unify with e_t.
    -- This ensures that 'x :> [1][]i32' does not make the second
    -- dimension unknown.  Use of matchDims is sensible because the
    -- structure of e_t' will be fully known due to the unification, and
    -- te_t because type expressions are complete.
    pure (te', toStruct te_t, e')
    where
      makeNonExtFresh ext = bitraverse onDim pure
        where
          onDim d@(Var v _ _)
            | qualLeaf v `elem` ext = pure d
          onDim d = do
            v <- newTypeName "coerce"
            constrain v . Size Nothing $
              mkUsage
                loc
                "a size coercion where the underlying expression size cannot be determined"
            pure $ sizeFromName (qualName v) (srclocOf 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 t = do
    constraints <- getConstraints
    -- These sizes will be immediately turned into existentials, so we
    -- do not need to care about their location.
    fst <$> sizeFree mempty (expKiller constraints) t
    where
      expKiller _ Var {} = Nothing
      expKiller constraints e =
        S.lookupMin $ S.filter (isUnknown constraints) $ (`S.difference` witnesses) $ fvVars $ freeInExp e
      isUnknown constraints vn
        | Just UnknownSize {} <- snd <$> M.lookup vn constraints = True
      isUnknown _ _ = False
      (witnesses, _) = determineSizeWitnesses $ toStruct t
  
  unscopeType ::
    SrcLoc ->
    [VName] ->
    TypeBase Size as ->
    TermTypeM (TypeBase Size as, [VName])
  unscopeType tloc unscoped =
    sizeFree tloc $ find (`elem` unscoped) . fvVars . freeInExp
  
  checkExp :: ExpBase NoInfo VName -> TermTypeM Exp
  checkExp (Literal val loc) =
    pure $ Literal val loc
  checkExp (Hole _ loc) = do
    t <- newTypeVar loc "t"
    pure $ Hole (Info t) loc
  checkExp (StringLit vs loc) =
    pure $ StringLit vs loc
  checkExp (IntLit val NoInfo loc) = do
    t <- newTypeVar loc "t"
    mustBeOneOf anyNumberType (mkUsage loc "integer literal") t
    pure $ IntLit val (Info t) loc
  checkExp (FloatLit val NoInfo loc) = do
    t <- newTypeVar loc "t"
    mustBeOneOf anyFloatType (mkUsage loc "float literal") t
    pure $ FloatLit val (Info t) loc
  checkExp (TupLit es loc) =
    TupLit <$> mapM checkExp es <*> pure loc
  checkExp (RecordLit fs loc) =
    RecordLit <$> evalStateT (mapM checkField fs) mempty <*> pure loc
    where
      checkField (RecordFieldExplicit f e rloc) = do
        errIfAlreadySet (unLoc f) rloc
        modify $ M.insert (unLoc f) rloc
        RecordFieldExplicit f <$> lift (checkExp e) <*> pure rloc
      checkField (RecordFieldImplicit name NoInfo rloc) = do
        errIfAlreadySet (baseName (unLoc name)) rloc
        t <- lift $ lookupVar rloc $ qualName $ unLoc name
        modify $ M.insert (baseName (unLoc name)) rloc
        pure $ RecordFieldImplicit name (Info t) rloc
  
      errIfAlreadySet f rloc = do
        maybe_sloc <- gets $ M.lookup f
        case maybe_sloc of
          Just sloc ->
            lift . typeError rloc mempty $
              "Field"
                <+> dquotes (pretty f)
                <+> "previously defined at"
                <+> pretty (locStrRel rloc sloc)
                <> "."
          Nothing -> 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 vs t loc) =
    pure $ ArrayVal vs t loc
  checkExp (ArrayLit all_es _ 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 all_es of
      [] -> do
        et <- newTypeVar loc "t"
        t <- arrayOfM loc et (Shape [sizeFromInteger 0 mempty])
        pure $ ArrayLit [] (Info t) loc
      e : es -> do
        e' <- checkExp e
        et <- expType e'
        es' <- mapM (unifies "type of first array element" et <=< checkExp) es
        t <- arrayOfM loc et (Shape [sizeFromInteger (genericLength all_es) mempty])
        pure $ ArrayLit (e' : es') (Info t) loc
  checkExp (AppExp (Range start maybe_step end loc) _) = do
    start' <- require "use in range expression" anySignedType =<< checkExp start
    start_t <- expType start'
    maybe_step' <- case maybe_step of
      Nothing -> pure Nothing
      Just step -> do
        let warning = warn loc "First and second element of range are identical, this will produce an empty array."
        case (start, step) of
          (Literal x _, Literal y _) -> when (x == y) warning
          (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning
          _ -> pure ()
        Just <$> (unifies "use in range expression" start_t =<< checkExp step)
  
    let unifyRange e = unifies "use in range expression" start_t =<< checkExp e
    end' <- traverse unifyRange end
  
    end_t <- case end' of
      DownToExclusive e -> expType e
      ToInclusive e -> expType e
      UpToExclusive e -> expType e
  
    -- Special case some ranges to give them a known size.
    let warnIfBinding binds size =
          if binds
            then do
              warn (srclocOf size) $
                withIndexLink
                  "size-expression-bind"
                  "Size expression with binding is replaced by unknown size."
              d <- newRigidDim loc RigidRange "range_dim"
              pure (sizeFromName (qualName d) mempty, Just d)
            else pure (size, Nothing)
    (dim, retext) <-
      case (isInt64 start', isInt64 <$> maybe_step', end') of
        (Just 0, Just (Just 1), UpToExclusive end'')
          | Scalar (Prim (Signed Int64)) <- end_t ->
              warnIfBinding (hasBinding end'') end''
        (Just 0, Nothing, UpToExclusive end'')
          | Scalar (Prim (Signed Int64)) <- end_t ->
              warnIfBinding (hasBinding end'') end''
        (_, Nothing, UpToExclusive end'')
          | Scalar (Prim (Signed Int64)) <- end_t ->
              warnIfBinding (hasBinding end'' || hasBinding start') $ sizeMinus end'' start'
        (_, Nothing, ToInclusive end'')
          -- No stride means we assume a stride of one.
          | Scalar (Prim (Signed Int64)) <- end_t ->
              warnIfBinding (hasBinding end'' || hasBinding start') $ sizeMinusInc end'' start'
        (Just 1, Just (Just 2), ToInclusive end'')
          | Scalar (Prim (Signed Int64)) <- end_t ->
              warnIfBinding (hasBinding end'') end''
        _ -> do
          d <- newRigidDim loc RigidRange "range_dim"
          pure (sizeFromName (qualName d) mempty, Just d)
  
    t <- arrayOfM loc start_t (Shape [dim])
    let res = AppRes t (maybeToList retext)
  
    pure $ AppExp (Range start' maybe_step' end' loc) (Info res)
    where
      i64 = Scalar $ Prim $ Signed Int64
      mkBinOp op t x y =
        AppExp
          ( BinOp
              (qualName (intrinsicVar op), mempty)
              sizeBinOpInfo
              (x, Info Nothing)
              (y, Info Nothing)
              mempty
          )
          (Info $ AppRes t [])
      mkSub = mkBinOp "-" i64
      mkAdd = mkBinOp "+" i64
      sizeMinus j i = j `mkSub` i
      sizeMinusInc j i = (j `mkSub` i) `mkAdd` sizeFromInteger 1 mempty
      sizeBinOpInfo = Info $ foldFunType [i64, i64] $ RetType [] i64
  checkExp (Ascript e te loc) = do
    (te', e') <- checkAscript loc te e
    pure $ Ascript e' te' loc
  checkExp (Coerce e te NoInfo loc) = do
    (te', te_t, e') <- checkCoerce loc te e
    t <- expTypeFully e'
    t' <- matchDims (const . const pure) t te_t
    pure $ Coerce e' te' (Info t') loc
  checkExp (AppExp (BinOp (op, oploc) NoInfo (e1, _) (e2, _) loc) NoInfo) = do
    ftype <- lookupVar oploc op
    e1' <- checkExp e1
    e2' <- checkExp e2
  
    -- Note that the application to the first operand cannot fix any
    -- existential sizes, because it must by necessity be a function.
    (_, rt, p1_ext, _) <- checkApply loc (Just op, 0) ftype e1'
    (_, rt', p2_ext, retext) <- checkApply loc (Just op, 1) rt e2'
  
    pure $
      AppExp
        ( BinOp
            (op, oploc)
            (Info ftype)
            (e1', Info p1_ext)
            (e2', Info p2_ext)
            loc
        )
        (Info (AppRes rt' retext))
  checkExp (Project k e NoInfo loc) = do
    e' <- checkExp e
    t <- expType e'
    kt <- mustHaveField (mkUsage loc $ docText $ "projection of field " <> dquotes (pretty k)) k t
    pure $ Project k e' (Info kt) loc
  checkExp (AppExp (If e1 e2 e3 loc) _) = do
    e1' <- checkExp e1
    e2' <- checkExp e2
    e3' <- checkExp e3
  
    let bool = Scalar $ Prim Bool
    e1_t <- expType e1'
    onFailure (CheckingRequired [bool] e1_t) $
      unify (mkUsage e1' "use as 'if' condition") bool e1_t
  
    (brancht, retext) <- unifyBranches loc e2' e3'
  
    zeroOrderType
      (mkUsage loc "returning value of this type from 'if' expression")
      "type returned from branch"
      brancht
  
    pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes brancht retext)
  checkExp (Parens e loc) =
    Parens <$> checkExp e <*> pure loc
  checkExp (QualParens (modname, modnameloc) e loc) = do
    mod <- lookupMod modname
    case mod of
      ModEnv env -> local (`withEnv` env) $ do
        e' <- checkExp e
        pure $ QualParens (modname, modnameloc) e' loc
      ModFun {} ->
        typeError loc mempty . withIndexLink "module-is-parametric" $
          "Module" <+> pretty modname <+> " is a parametric module."
  checkExp (Var qn NoInfo loc) = do
    t <- lookupVar loc qn
    pure $ Var qn (Info t) loc
  checkExp (Negate arg loc) = do
    arg' <- require "numeric negation" anyNumberType =<< checkExp arg
    pure $ Negate arg' loc
  checkExp (Not arg loc) = do
    arg' <- require "logical negation" (Bool : anyIntType) =<< checkExp arg
    pure $ Not arg' loc
  checkExp (AppExp (Apply fe args loc) NoInfo) = do
    fe' <- checkExp fe
    args' <- mapM (checkExp . snd) args
    t <- expType fe'
    let fname =
          case fe' of
            Var v _ _ -> Just v
            _ -> Nothing
    ((_, exts, rt), args'') <- mapAccumLM (onArg fname) (0, [], t) args'
  
    pure $ AppExp (Apply fe' args'' loc) $ Info $ AppRes rt exts
    where
      onArg fname (i, all_exts, t) arg' = do
        (_, rt, argext, exts) <- checkApply loc (fname, i) t arg'
        pure
          ( (i + 1, all_exts <> exts, rt),
            (Info argext, arg')
          )
  checkExp (AppExp (LetPat sizes pat e body loc) _) = do
    e' <- checkExp e
  
    -- Not technically an ascription, but we want the pattern to have
    -- exactly the type of 'e'.
    t <- expType e'
    bindingSizes sizes . incLevel . bindingPat sizes pat t $ \pat' -> do
      body' <- incLevel $ checkExp body
      body_t <- expTypeFully body'
  
      -- If the bound expression is of type i64, then we replace the
      -- pattern name with the expression in the type of the body.
      -- Otherwise, we need to come up with unknown sizes for the
      -- sizes going out of scope.
      t' <- normType t -- Might be overloaded integer until now.
      (body_t', retext) <-
        case (t', patNames pat') of
          (Scalar (Prim (Signed Int64)), [v])
            | not $ hasBinding e' -> do
                let f x = if x == v then Just (ExpSubst e') else Nothing
                pure (applySubst f body_t, [])
          _ ->
            unscopeType loc (map sizeName sizes <> patNames pat') body_t
  
      pure $
        AppExp
          (LetPat sizes (fmap toStruct pat') e' body' loc)
          (Info $ AppRes body_t' retext)
  checkExp (AppExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body loc) _) = do
    (tparams', params', maybe_retdecl', rettype, e') <-
      checkBinding (fst name, maybe_retdecl, tparams, params, e, loc)
  
    let entry = BoundV tparams' $ funType params' rettype
        bindF scope =
          scope
            { scopeVtable = M.insert (fst name) entry $ scopeVtable scope
            }
    body' <- localScope bindF $ checkExp body
  
    (body_t, ext) <- unscopeType loc [fst name] =<< expTypeFully body'
  
    pure $
      AppExp
        ( LetFun
            name
            (tparams', params', maybe_retdecl', Info rettype, e')
            body'
            loc
        )
        (Info $ AppRes body_t ext)
  checkExp (AppExp (LetWith dest src slice ve body loc) _) = do
    src' <- checkIdent src
    slice' <- checkSlice slice
    (t, _) <- newArrayType (mkUsage src "type of source array") "src" $ sliceDims slice'
    unify (mkUsage loc "type of target array") t $ unInfo $ identType src'
  
    (elemt, _) <- sliceShape (Just (loc, Nonrigid)) slice' =<< normTypeFully t
  
    ve' <- unifies "type of target array" elemt =<< checkExp ve
  
    bindingIdent dest (unInfo (identType src')) $ \dest' -> do
      body' <- checkExp body
      (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'
      pure $ AppExp (LetWith dest' src' slice' ve' body' loc) (Info $ AppRes body_t ext)
  checkExp (Update src slice ve loc) = do
    slice' <- checkSlice slice
    (t, _) <- newArrayType (mkUsage' src) "src" $ sliceDims slice'
    (elemt, _) <- sliceShape (Just (loc, Nonrigid)) slice' =<< normTypeFully t
    ve' <- unifies "type of target array" elemt =<< checkExp ve
    src' <- unifies "type of target array" t =<< checkExp src
    pure $ Update src' slice' ve' 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 src fields ve NoInfo loc) = do
    src' <- checkExp src
    ve' <- checkExp ve
    a <- expTypeFully src'
    foldM_ (flip $ mustHaveField usage) a fields
    ve_t <- expType ve'
    updated_t <- updateField fields ve_t =<< expTypeFully src'
    pure $ RecordUpdate src' fields ve' (Info updated_t) loc
    where
      usage = mkUsage loc "record update"
      updateField [] ve_t src_t = do
        (src_t', _) <- allDimsFreshInType usage Nonrigid "any" src_t
        onFailure (CheckingRecordUpdate fields src_t' ve_t) $
          unify usage src_t' ve_t
        pure ve_t
      updateField (f : fs) ve_t (Scalar (Record m))
        | Just f_t <- M.lookup f m = do
            f_t' <- updateField fs ve_t f_t
            pure $ Scalar $ Record $ M.insert f f_t' m
      updateField _ _ _ =
        typeError loc mempty . withIndexLink "record-type-not-known" $
          "Full type of"
            </> indent 2 (pretty src)
            </> textwrap " is not known at this point.  Add a type annotation to the original record to disambiguate."
  
  --
  checkExp (AppExp (Index e slice loc) _) = do
    slice' <- checkSlice slice
    (t, _) <- newArrayType (mkUsage' loc) "e" $ sliceDims slice'
    e' <- unifies "being indexed at" t =<< checkExp e
    -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value.
    (t', retext) <-
      sliceShape (Just (loc, Rigid (RigidSlice Nothing ""))) slice'
        =<< expTypeFully e'
  
    pure $ AppExp (Index e' slice' loc) (Info $ AppRes t' retext)
  checkExp (Assert e1 e2 NoInfo loc) = do
    e1' <- require "being asserted" [Bool] =<< checkExp e1
    e2' <- checkExp e2
    pure $ Assert e1' e2' (Info (prettyText e1)) loc
  checkExp (Lambda params body rettype_te NoInfo loc) = do
    (params', body', rettype', RetType dims ty) <-
      incLevel . bindingParams [] params $ \params' -> do
        rettype_checked <- traverse checkTypeExpNonrigid rettype_te
        let declared_rettype =
              case rettype_checked of
                Just (_, st, _) -> Just st
                Nothing -> Nothing
        body' <- checkFunBody params' body declared_rettype loc
        body_t <- expTypeFully body'
  
        params'' <- mapM updateTypes params'
  
        (rettype', rettype_st) <-
          case rettype_checked of
            Just (te, st, ext) ->
              pure (Just te, RetType ext st)
            Nothing -> do
              ret <- inferReturnSizes params'' $ toRes Nonunique body_t
              pure (Nothing, ret)
  
        pure (params'', body', rettype', rettype_st)
  
    verifyFunctionParams Nothing params'
  
    (ty', dims') <- unscopeType loc dims ty
  
    pure $ Lambda params' body' rettype' (Info (RetType dims' ty')) loc
    where
      -- Inferring the sizes of the return type of a lambda is a lot
      -- like let-generalisation.  We wish to remove any rigid sizes
      -- that were created when checking the body, except for those that
      -- are visible in types that existed before we entered the body,
      -- are parameters, or are used in parameters.
      inferReturnSizes params' ret = do
        cur_lvl <- curLevel
        let named (Named x, _, _) = Just x
            named (Unnamed, _, _) = Nothing
            param_names = mapMaybe (named . patternParam) params'
            pos_sizes =
              sizeNamesPos $ funType params' $ RetType [] ret
            hide k (lvl, _) =
              lvl >= cur_lvl && k `notElem` param_names && k `S.notMember` pos_sizes
  
        hidden_sizes <-
          S.fromList . M.keys . M.filterWithKey hide <$> getConstraints
  
        let onDim name
              | name `S.member` hidden_sizes = S.singleton name
            onDim _ = mempty
  
        pure $ RetType (S.toList $ foldMap onDim $ fvVars $ freeInType ret) ret
  checkExp (OpSection op _ loc) = do
    ftype <- lookupVar loc op
    pure $ OpSection op (Info ftype) loc
  checkExp (OpSectionLeft op _ e _ _ loc) = do
    ftype <- lookupVar loc op
    e' <- checkExp e
    (t1, rt, argext, retext) <- checkApply loc (Just op, 0) ftype e'
    case (ftype, rt) of
      (Scalar (Arrow _ m1 d1 _ _), Scalar (Arrow _ m2 d2 t2 rettype)) ->
        pure $
          OpSectionLeft
            op
            (Info ftype)
            e'
            (Info (m1, toParam d1 t1, argext), Info (m2, toParam d2 t2))
            (Info rettype, Info retext)
            loc
      _ ->
        typeError loc mempty $
          "Operator section with invalid operator of type" <+> pretty ftype
  checkExp (OpSectionRight op _ e _ NoInfo loc) = do
    ftype <- lookupVar loc op
    e' <- checkExp e
    case ftype of
      Scalar (Arrow _ m1 d1 t1 (RetType [] (Scalar (Arrow _ m2 d2 t2 (RetType dims2 ret))))) -> do
        (t2', arrow', argext, _) <-
          checkApply
            loc
            (Just op, 1)
            (Scalar $ Arrow mempty m2 d2 t2 $ RetType [] $ Scalar $ Arrow Nonunique m1 d1 t1 $ RetType dims2 ret)
            e'
        case arrow' of
          Scalar (Arrow _ _ _ t1' (RetType dims2' ret')) ->
            pure $
              OpSectionRight
                op
                (Info ftype)
                e'
                (Info (m1, toParam d1 t1'), Info (m2, toParam d2 t2', argext))
                (Info $ RetType dims2' ret')
                loc
          _ -> error $ "OpSectionRight: impossible type\n" <> prettyString arrow'
      _ ->
        typeError loc mempty $
          "Operator section with invalid operator of type" <+> pretty ftype
  checkExp (ProjectSection fields NoInfo loc) = do
    a <- newTypeVar loc "a"
    let usage = mkUsage loc "projection at"
    b <- foldM (flip $ mustHaveField usage) a fields
    let ft = Scalar $ Arrow mempty Unnamed Observe a $ RetType [] $ toRes Nonunique b
    pure $ ProjectSection fields (Info ft) loc
  checkExp (IndexSection slice NoInfo loc) = do
    slice' <- checkSlice slice
    (t, _) <- newArrayType (mkUsage' loc) "e" $ sliceDims slice'
    (t', retext) <- sliceShape Nothing slice' t
    let ft = Scalar $ Arrow mempty Unnamed Observe t $ RetType retext $ toRes Nonunique t'
    pure $ IndexSection slice' (Info ft) loc
  checkExp (AppExp (Loop _ mergepat loopinit form loopbody loc) _) = do
    ((sparams, mergepat', loopinit', form', loopbody'), appres) <-
      checkLoop checkExp (mergepat, loopinit, form, loopbody) loc
    pure $
      AppExp
        (Loop sparams mergepat' loopinit' form' loopbody' loc)
        (Info appres)
  checkExp (Constr name es NoInfo loc) = do
    t <- newTypeVar loc "t"
    es' <- mapM checkExp es
    ets <- mapM expType es'
    mustHaveConstr (mkUsage loc "use of constructor") name t ets
    pure $ Constr name es' (Info t) loc
  checkExp (AppExp (Match e cs loc) _) = do
    e' <- checkExp e
    mt <- expType e'
    (cs', t, retext) <- checkCases mt cs
    zeroOrderType
      (mkUsage loc "being returned 'match'")
      "type returned from pattern match"
      t
    pure $ AppExp (Match e' cs' loc) (Info $ AppRes t retext)
  checkExp (Attr info e loc) =
    Attr <$> checkAttr info <*> checkExp e <*> pure loc
  
  checkCases ::
    StructType ->
    NE.NonEmpty (CaseBase NoInfo VName) ->
    TermTypeM (NE.NonEmpty (CaseBase Info VName), StructType, [VName])
  checkCases mt rest_cs =
    case NE.uncons rest_cs of
      (c, Nothing) -> do
        (c', t, retext) <- checkCase mt c
        pure (NE.singleton c', t, retext)
      (c, Just cs) -> do
        ((c', c_t, _), (cs', cs_t, _)) <-
          (,) <$> checkCase mt c <*> checkCases mt cs
        (brancht, retext) <- unifyBranchTypes (srclocOf c) c_t cs_t
        pure (NE.cons c' cs', brancht, retext)
  
  checkCase ::
    StructType ->
    CaseBase NoInfo VName ->
    TermTypeM (CaseBase Info VName, StructType, [VName])
  checkCase mt (CasePat p e loc) =
    bindingPat [] p mt $ \p' -> do
      e' <- checkExp e
      e_t <- expTypeFully e'
      (e_t', retext) <- unscopeType loc (patNames p') e_t
      pure (CasePat (fmap toStruct p') e' loc, e_t', 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 (Functor, Show)
  
  instance Pretty (Unmatched (Pat StructType)) where
    pretty um = case um of
      (UnmatchedNum p nums) -> pretty' p <+> "where p is not one of" <+> pretty nums
      (UnmatchedBool p) -> pretty' p
      (UnmatchedConstr p) -> pretty' p
      (Unmatched p) -> pretty' p
      where
        pretty' (PatAscription p t _) = pretty p <> ":" <+> pretty t
        pretty' (PatParens p _) = parens $ pretty' p
        pretty' (PatAttr _ p _) = parens $ pretty' p
        pretty' (Id v _ _) = prettyName v
        pretty' (TuplePat pats _) = parens $ commasep $ map pretty' pats
        pretty' (RecordPat fs _) = braces $ commasep $ map ppField fs
          where
            ppField (L _ name, t) = pretty (nameToString name) <> equals <> pretty' t
        pretty' Wildcard {} = "_"
        pretty' (PatLit e _ _) = pretty e
        pretty' (PatConstr n _ ps _) = "#" <> pretty n <+> sep (map pretty' ps)
  
  checkIdent :: IdentBase NoInfo VName StructType -> TermTypeM (Ident StructType)
  checkIdent (Ident name _ loc) = do
    vt <- lookupVar loc $ qualName name
    pure $ Ident name (Info vt) loc
  
  checkSlice :: SliceBase NoInfo VName -> TermTypeM [DimIndex]
  checkSlice = mapM checkDimIndex
    where
      checkDimIndex (DimFix i) = do
        DimFix <$> (require "use as index" anySignedType =<< checkExp i)
      checkDimIndex (DimSlice i j s) =
        DimSlice <$> check i <*> check j <*> check s
  
      check =
        maybe (pure Nothing) $
          fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int64) <=< checkExp
  
  -- The number of dimensions affected by this slice (so the minimum
  -- rank of the array we are slicing).
  sliceDims :: [DimIndex] -> Int
  sliceDims = length
  
  instantiateDimsInReturnType ::
    SrcLoc ->
    Maybe (QualName VName) ->
    ResRetType ->
    TermTypeM (ResType, [VName])
  instantiateDimsInReturnType loc fname (RetType dims t)
    | null dims =
        pure (t, mempty)
    | otherwise = do
        dims' <- mapM new dims
        pure (first (onDim $ zip dims $ map (ExpSubst . (`sizeFromName` loc) . qualName) dims') t, dims')
    where
      new =
        newRigidDim loc (RigidRet fname)
          . nameFromText
          . T.takeWhile isAscii
          . baseText
      onDim dims' = applySubst (`lookup` 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 (Array _ _ t) = boundInsideType (Scalar t)
  boundInsideType (Scalar Prim {}) = mempty
  boundInsideType (Scalar (TypeVar _ _ targs)) = foldMap f targs
    where
      f (TypeArgType t) = boundInsideType t
      f TypeArgDim {} = mempty
  boundInsideType (Scalar (Record fs)) = foldMap boundInsideType fs
  boundInsideType (Scalar (Sum cs)) = foldMap (foldMap boundInsideType) cs
  boundInsideType (Scalar (Arrow _ pn _ t1 (RetType dims t2))) =
    pn' <> boundInsideType t1 <> S.fromList dims <> boundInsideType t2
    where
      pn' = case pn of
        Unnamed -> mempty
        Named v -> S.singleton 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 = flip execState mempty . traverseDims f
    where
      f bound pos e =
        case pos of
          PosImmediate ->
            modify ((fvVars fv, mempty) <>)
          PosParam ->
            modify ((mempty, fvVars fv) <>)
          PosReturn -> pure ()
        where
          fv = freeInExp e `freeWithout` bound
  
  checkApply ::
    SrcLoc ->
    ApplyOp ->
    StructType ->
    Exp ->
    TermTypeM (StructType, StructType, Maybe VName, [VName])
  checkApply loc (fname, _) (Scalar (Arrow _ pname _ tp1 tp2)) argexp = do
    let argtype = typeOf argexp
    onFailure (CheckingApply fname argexp tp1 argtype) $ do
      unify (mkUsage argexp "use as function argument") tp1 argtype
  
      -- Perform substitutions of instantiated variables in the types.
      (tp2', ext) <- instantiateDimsInReturnType loc fname =<< normTypeFully tp2
      argtype' <- normTypeFully argtype
  
      -- Check whether this would produce an impossible return type.
      let (tp2_produced_dims, tp2_paramdims) = dimUses tp2'
          problematic = S.fromList ext <> boundInsideType argtype'
          problem = any (`S.member` problematic) (tp2_paramdims `S.difference` tp2_produced_dims)
      when (not (S.null problematic) && problem) $ do
        typeError loc mempty . withIndexLink "existential-param-ret" $
          "Existential size would appear in function parameter of return type:"
            </> indent 2 (pretty (RetType ext tp2'))
            </> textwrap "This is usually because a higher-order function is used with functional arguments that return existential sizes or locally named sizes, which are then used as parameters of other function arguments."
  
      (argext, tp2'') <-
        case pname of
          Named pname'
            | S.member pname' (fvVars $ freeInType tp2') ->
                if hasBinding argexp
                  then do
                    warn (srclocOf argexp) $
                      withIndexLink
                        "size-expression-bind"
                        "Size expression with binding is replaced by unknown size."
                    d <- newRigidDim argexp (RigidArg fname $ prettyTextOneLine $ bareExp argexp) "n"
                    let parsubst v =
                          if v == pname'
                            then Just $ ExpSubst $ sizeFromName (qualName d) $ srclocOf argexp
                            else Nothing
                    pure (Just d, applySubst parsubst $ toStruct tp2')
                  else
                    let parsubst v =
                          if v == pname'
                            then Just $ ExpSubst $ fromMaybe argexp $ stripExp argexp
                            else Nothing
                     in pure (Nothing, applySubst parsubst $ toStruct tp2')
          _ -> pure (Nothing, toStruct tp2')
  
      pure (tp1, tp2'', argext, ext)
  checkApply loc fname tfun@(Scalar TypeVar {}) arg = do
    tv <- newTypeVar loc "b"
    unify (mkUsage loc "use as function") tfun $
      Scalar (Arrow mempty Unnamed Observe (typeOf arg) $ RetType [] $ paramToRes tv)
    tfun' <- normType tfun
    checkApply loc fname tfun' arg
  checkApply loc (fname, prev_applied) ftype argexp = do
    let fname' = maybe "expression" (dquotes . pretty) fname
  
    typeError loc mempty $
      if prev_applied == 0
        then
          "Cannot apply"
            <+> fname'
            <+> "as function, as it has type:"
            </> indent 2 (pretty ftype)
        else
          "Cannot apply"
            <+> fname'
            <+> "to argument #"
            <> pretty (prev_applied + 1)
              <+> dquotes (shorten $ group $ pretty argexp)
            <> ","
              </> "as"
              <+> fname'
              <+> "only takes"
              <+> pretty prev_applied
              <+> arguments
            <> "."
    where
      arguments
        | prev_applied == 1 = "argument"
        | otherwise = "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 e = runTermTypeM checkExp $ do
    e' <- checkExp e
    (tparams, _, RetType _ t') <-
      letGeneralise (nameFromString "<exp>") (srclocOf e) [] [] $
        toRes Nonunique $
          typeOf e'
    fixOverloadedTypes $ typeVars t'
    e'' <- normTypeFully e'
    localChecks e''
    causalityCheck e''
    pure (tparams, e'')
  
  -- | Type-check a single size expression in isolation.  This expression may
  -- turn out to be polymorphic, in which case it is unified with i64.
  checkSizeExp :: ExpBase NoInfo VName -> TypeM Exp
  checkSizeExp e = runTermTypeM checkExp $ do
    e' <- checkExp e
    let t = typeOf e'
    when (hasBinding e') $
      typeError (srclocOf e') mempty . withIndexLink "size-expression-bind" $
        "Size expression with binding is forbidden."
    unify (mkUsage e' "Size expression") t (Scalar (Prim (Signed Int64)))
    normTypeFully e'
  
  -- Verify that all sum type constructors and empty array literals have
  -- a size that is known (rigid or a type parameter).  This is to
  -- ensure that we can actually determine their shape at run-time.
  causalityCheck :: Exp -> TermTypeM ()
  causalityCheck binding_body = do
    constraints <- getConstraints
  
    let checkCausality what known t loc
          | (d, dloc) : _ <-
              mapMaybe (unknown constraints known) $
                S.toList (fvVars $ freeInType t) =
              Just $ lift $ causality what (locOf loc) d dloc t
          | otherwise = Nothing
  
        checkParamCausality known p =
          checkCausality (pretty p) known (patternType p) (locOf p)
  
        collectingNewKnown = lift . flip execStateT mempty
  
        onExp ::
          S.Set VName ->
          Exp ->
          StateT (S.Set VName) (Either TypeError) Exp
  
        onExp known (Var v (Info t) loc)
          | Just bad <- checkCausality (dquotes (pretty v)) known t loc =
              bad
        onExp known (ProjectSection _ (Info t) loc)
          | Just bad <- checkCausality "projection section" known t loc =
              bad
        onExp known (IndexSection _ (Info t) loc)
          | Just bad <- checkCausality "projection section" known t loc =
              bad
        onExp known (OpSectionRight _ (Info t) _ _ _ loc)
          | Just bad <- checkCausality "operator section" known t loc =
              bad
        onExp known (OpSectionLeft _ (Info t) _ _ _ loc)
          | Just bad <- checkCausality "operator section" known t loc =
              bad
        onExp known (ArrayLit [] (Info t) loc)
          | Just bad <- checkCausality "empty array" known t loc =
              bad
        onExp known (Hole (Info t) loc)
          | Just bad <- checkCausality "hole" known t loc =
              bad
        onExp known e@(Lambda params body _ _ _)
          | bad : _ <- mapMaybe (checkParamCausality known) params =
              bad
          | otherwise = do
              -- Existentials coming into existence in the lambda body
              -- are not known outside of it.
              void $ collectingNewKnown $ onExp known body
              pure e
        onExp known e@(AppExp (LetPat _ _ bindee_e body_e _) (Info res)) = do
          sequencePoint known bindee_e body_e $ appResExt res
          pure e
        onExp known e@(AppExp (Match scrutinee cs _) (Info res)) = do
          new_known <- collectingNewKnown $ onExp known scrutinee
          void $ recurse (new_known <> known) cs
          modify ((new_known <> S.fromList (appResExt res)) <>)
          pure e
        onExp known e@(AppExp (Apply f args _) (Info res)) = do
          seqArgs known $ reverse $ NE.toList args
          pure e
          where
            seqArgs known' [] = do
              void $ onExp known' f
              modify (S.fromList (appResExt res) <>)
            seqArgs known' ((Info p, x) : xs) = do
              new_known <- collectingNewKnown $ onExp known' x
              void $ seqArgs (new_known <> known') xs
              modify ((new_known <> S.fromList (maybeToList p)) <>)
        onExp known e@(Constr v args (Info t) loc) = do
          seqArgs known args
          pure e
          where
            seqArgs known' []
              | Just bad <- checkCausality (dquotes ("#" <> pretty v)) known' t loc =
                  bad
              | otherwise =
                  pure ()
            seqArgs known' (x : xs) = do
              new_known <- collectingNewKnown $ onExp known' x
              void $ seqArgs (new_known <> known') xs
              modify (new_known <>)
        onExp
          known
          e@(AppExp (BinOp (f, floc) ft (x, Info xp) (y, Info yp) _) (Info res)) = do
            args_known <-
              collectingNewKnown $ sequencePoint known x y $ catMaybes [xp, yp]
            void $ onExp (args_known <> known) (Var f ft floc)
            modify ((args_known <> S.fromList (appResExt res)) <>)
            pure e
        onExp known e@(AppExp e' (Info res)) = do
          recurse known e'
          modify (<> S.fromList (appResExt res))
          pure e
        onExp known e = do
          recurse known e
          pure e
  
        recurse known = void . astMap mapper
          where
            mapper = identityMapper {mapOnExp = onExp known}
  
        sequencePoint known x y ext = do
          new_known <- collectingNewKnown $ onExp known x
          void $ onExp (new_known <> known) y
          modify ((new_known <> S.fromList ext) <>)
  
    either throwError (const $ pure ()) $
      evalStateT (onExp mempty binding_body) mempty
    where
      unknown constraints known v = do
        guard $ v `S.notMember` known
        loc <- case snd <$> M.lookup v constraints of
          Just (UnknownSize loc _) -> Just loc
          _ -> Nothing
        pure (v, loc)
  
      causality what loc d dloc t =
        Left . TypeError loc mempty . withIndexLink "causality-check" $
          "Causality check: size"
            <+> dquotes (prettyName d)
            <+> "needed for type of"
            <+> what
            <> colon
              </> indent 2 (pretty t)
              </> "But"
              <+> dquotes (prettyName d)
              <+> "is computed at"
              <+> pretty (locStrRel loc dloc)
            <> "."
              </> ""
              </> "Hint:"
              <+> align
                ( textwrap "Bind the expression producing"
                    <+> dquotes (prettyName d)
                    <+> "with 'let' beforehand."
                )
  
  mustBeIrrefutable :: (MonadTypeChecker f) => Pat StructType -> f ()
  mustBeIrrefutable p = do
    case unmatched [p] of
      [] -> pure ()
      ps' ->
        typeError p mempty . withIndexLink "refutable-pattern" $
          "Refutable pattern not allowed here.\nUnmatched cases:"
            </> indent 2 (stack (map pretty 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 = void . check
    where
      check e@(AppExp (Match _ cs loc) _) = do
        let ps = fmap (\(CasePat p _ _) -> p) cs
        case unmatched $ NE.toList ps of
          [] -> recurse e
          ps' ->
            typeError loc mempty . withIndexLink "unmatched-cases" $
              "Unmatched cases in match expression:"
                </> indent 2 (stack (map pretty ps'))
      check e@(AppExp (LetPat _ p _ _ _) _) =
        mustBeIrrefutable p *> recurse e
      check e@(Lambda ps _ _ _ _) =
        mapM_ (mustBeIrrefutable . fmap toStruct) ps *> recurse e
      check e@(AppExp (LetFun _ (_, ps, _, _, _) _ _) _) =
        mapM_ (mustBeIrrefutable . fmap toStruct) ps *> recurse e
      check e@(AppExp (Loop _ p _ form _ _) _) = do
        mustBeIrrefutable (fmap toStruct p)
        case form of
          ForIn form_p _ -> mustBeIrrefutable form_p
          _ -> pure ()
        recurse e
      check e@(IntLit x ty loc) =
        e <$ case ty of
          Info (Scalar (Prim t)) -> errorBounds (inBoundsI x t) x t loc
          _ -> error "Inferred type of int literal is not a number"
      check e@(FloatLit x ty loc) =
        e <$ case ty of
          Info (Scalar (Prim (FloatType t))) -> errorBounds (inBoundsF x t) x t loc
          _ -> error "Inferred type of float literal is not a float"
      check e@(Negate (IntLit x ty loc1) loc2) =
        e <$ case ty of
          Info (Scalar (Prim t)) -> errorBounds (inBoundsI (-x) t) (-x) t (loc1 <> loc2)
          _ -> error "Inferred type of int literal is not a number"
      check e@(AppExp (BinOp (QualName [] v, _) _ (x, _) _ loc) _)
        | baseName v == "==",
          Array {} <- typeOf x,
          baseTag v <= maxIntrinsicTag = do
            warn loc $
              textwrap
                "Comparing arrays with \"==\" is deprecated and will stop working in a future revision of the language."
            recurse e
      check e = recurse e
      recurse = astMap identityMapper {mapOnExp = check}
  
      bitWidth ty = 8 * intByteSize ty :: Int
  
      inBoundsI x (Signed t) = x >= -2 ^ (bitWidth t - 1) && x < 2 ^ (bitWidth t - 1)
      inBoundsI x (Unsigned t) = x >= 0 && x < 2 ^ bitWidth t
      inBoundsI x (FloatType Float16) = not $ isInfinite (fromIntegral x :: Half)
      inBoundsI x (FloatType Float32) = not $ isInfinite (fromIntegral x :: Float)
      inBoundsI x (FloatType Float64) = not $ isInfinite (fromIntegral x :: Double)
      inBoundsI _ Bool = error "Inferred type of int literal is not a number"
      inBoundsF x Float16 = not $ isInfinite (realToFrac x :: Float)
      inBoundsF x Float32 = not $ isInfinite (realToFrac x :: Float)
      inBoundsF x Float64 = not $ isInfinite x
  
      errorBounds inBounds x ty loc =
        unless inBounds $
          typeError loc mempty . withIndexLink "literal-out-of-bounds" $
            "Literal "
              <> pretty x
              <> " out of bounds for inferred type "
              <> pretty ty
              <> "."
  
  -- | 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 (fname, maybe_retdecl, tparams, params, body, loc) =
    runTermTypeM checkExp $ do
      (tparams', params', maybe_retdecl', RetType dims rettype', body') <-
        checkBinding (fname, maybe_retdecl, tparams, params, body, loc)
  
      -- Since this is a top-level function, we also resolve overloaded
      -- types, using either defaults or complaining about ambiguities.
      fixOverloadedTypes $
        typeVars rettype' <> foldMap (typeVars . patternType) params'
  
      -- Then replace all inferred types in the body and parameters.
      body'' <- normTypeFully body'
      params'' <- mapM normTypeFully params'
      maybe_retdecl'' <- traverse updateTypes maybe_retdecl'
      rettype'' <- normTypeFully rettype'
  
      -- Check if the function body can actually be evaluated.
      causalityCheck body''
  
      -- Check for various problems.
      mapM_ (mustBeIrrefutable . fmap toStruct) params'
      localChecks body''
  
      let ((body''', updated_ret), errors) =
            Consumption.checkValDef
              ( fname,
                params'',
                body'',
                RetType dims rettype'',
                maybe_retdecl'',
                loc
              )
  
      mapM_ throwError errors
  
      pure (tparams', params'', maybe_retdecl'', updated_ret, body''')
  
  -- | This is "fixing" as in "setting them", not "correcting them".  We
  -- only make very conservative fixing.
  fixOverloadedTypes :: Names -> TermTypeM ()
  fixOverloadedTypes tyvars_at_toplevel =
    getConstraints >>= mapM_ fixOverloaded . M.toList . M.map snd
    where
      fixOverloaded (v, Overloaded ots usage)
        | Signed Int32 `elem` ots = do
            unify usage (Scalar (TypeVar mempty (qualName v) [])) $
              Scalar (Prim $ Signed Int32)
            when (v `S.member` tyvars_at_toplevel) $
              warn usage "Defaulting ambiguous type to i32."
        | FloatType Float64 `elem` ots = do
            unify usage (Scalar (TypeVar mempty (qualName v) [])) $
              Scalar (Prim $ FloatType Float64)
            when (v `S.member` tyvars_at_toplevel) $
              warn usage "Defaulting ambiguous type to f64."
        | otherwise =
            typeError usage mempty . withIndexLink "ambiguous-type" $
              "Type is ambiguous (could be one of"
                <+> commasep (map pretty ots)
                <> ")."
                  </> "Add a type annotation to disambiguate the type."
      fixOverloaded (v, NoConstraint _ usage) = do
        -- See #1552.
        unify usage (Scalar (TypeVar mempty (qualName v) [])) $
          Scalar (tupleRecord [])
        when (v `S.member` tyvars_at_toplevel) $
          warn usage "Defaulting ambiguous type to ()."
      fixOverloaded (_, Equality usage) =
        typeError usage mempty . withIndexLink "ambiguous-type" $
          "Type is ambiguous (must be equality type)."
            </> "Add a type annotation to disambiguate the type."
      fixOverloaded (_, HasFields _ fs usage) =
        typeError usage mempty . withIndexLink "ambiguous-type" $
          "Type is ambiguous.  Must be record with fields:"
            </> indent 2 (stack $ map field $ M.toList fs)
            </> "Add a type annotation to disambiguate the type."
        where
          field (l, t) = pretty l <> colon <+> align (pretty t)
      fixOverloaded (_, HasConstrs _ cs usage) =
        typeError usage mempty . withIndexLink "ambiguous-type" $
          "Type is ambiguous (must be a sum type with constructors:"
            <+> pretty (Sum cs)
            <> ")."
              </> "Add a type annotation to disambiguate the type."
      fixOverloaded (v, Size Nothing (Usage Nothing loc)) =
        typeError loc mempty . withIndexLink "ambiguous-size" $
          "Ambiguous size" <+> dquotes (prettyName v) <> "."
      fixOverloaded (v, Size Nothing (Usage (Just u) loc)) =
        typeError loc mempty . withIndexLink "ambiguous-size" $
          "Ambiguous size" <+> dquotes (prettyName v) <+> "arising from" <+> pretty u <> "."
      fixOverloaded _ = pure ()
  
  hiddenParamNames :: [Pat ParamType] -> [VName]
  hiddenParamNames params = hidden
    where
      param_all_names = mconcat $ map patNames params
      named (Named x, _, _) = Just x
      named (Unnamed, _, _) = Nothing
      param_names =
        S.fromList $ mapMaybe (named . patternParam) params
      hidden = filter (`notElem` param_names) param_all_names
  
  inferredReturnType :: SrcLoc -> [Pat ParamType] -> StructType -> TermTypeM StructType
  inferredReturnType loc params 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.
    fst <$> unscopeType loc hidden_params t
    where
      hidden_params = filter (`elem` hidden) $ foldMap patNames params
      hidden = hiddenParamNames 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 (fname, maybe_retdecl, tparams, params, body, loc) =
    incLevel . bindingParams tparams params $ \params' -> do
      maybe_retdecl' <- traverse checkTypeExpNonrigid maybe_retdecl
  
      body' <-
        checkFunBody
          params'
          body
          ((\(_, x, _) -> x) <$> maybe_retdecl')
          (maybe loc srclocOf maybe_retdecl)
  
      params'' <- mapM updateTypes params'
      body_t <- expTypeFully body'
  
      (maybe_retdecl'', rettype) <- case maybe_retdecl' of
        Just (retdecl', ret, _) -> do
          ret' <- normTypeFully ret
          pure (Just retdecl', ret')
        Nothing
          | null params ->
              pure (Nothing, toRes Nonunique body_t)
          | otherwise -> do
              body_t' <- inferredReturnType loc params'' body_t
              pure (Nothing, toRes Nonunique body_t')
  
      verifyFunctionParams (Just fname) params''
  
      (tparams', params''', rettype') <-
        letGeneralise (baseName fname) loc tparams params'' =<< unscopeUnknown rettype
  
      when
        ( null params
            && any isSizeParam tparams'
            && not (null (retDims rettype'))
        )
        $ typeError loc mempty
        $ textwrap "A size-polymorphic value binding may not have a type with an existential size."
          </> "Type of this binding is:"
          </> indent 2 (pretty rettype')
          </> "with the following type parameters:"
          </> indent 2 (sep $ map pretty $ filter isSizeParam tparams')
  
      pure (tparams', params''', maybe_retdecl'', rettype', body')
  
  -- | Extract all the shape names that occur in positive position
  -- (roughly, left side of an arrow) in a given type.
  sizeNamesPos :: TypeBase Size als -> S.Set VName
  sizeNamesPos (Scalar (Arrow _ _ _ t1 (RetType _ t2))) = onParam t1 <> sizeNamesPos t2
    where
      onParam :: TypeBase Size als -> S.Set VName
      onParam (Scalar Arrow {}) = mempty
      onParam (Scalar (Record fs)) = mconcat $ map onParam $ M.elems fs
      onParam (Scalar (TypeVar _ _ targs)) = mconcat $ map onTypeArg targs
      onParam t = fvVars $ freeInType t
      onTypeArg (TypeArgDim (Var d _ _)) = S.singleton $ qualLeaf d
      onTypeArg (TypeArgDim _) = mempty
      onTypeArg (TypeArgType t) = onParam t
  sizeNamesPos _ = 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 fname params =
    onFailure (CheckingParams (baseName <$> fname)) $
      verifyParams (foldMap patNames params) =<< mapM updateTypes params
    where
      verifyParams forbidden (p : ps)
        | d : _ <- filter (`elem` forbidden) $ S.toList $ fvVars $ freeInPat p =
            typeError p mempty . withIndexLink "inaccessible-size" $
              "Parameter"
                <+> dquotes (pretty p)
                </> "refers to size"
                <+> dquotes (prettyName d)
                <> comma
                  </> textwrap "which will not be accessible to the caller"
                <> comma
                  </> textwrap "possibly because it is nested in a tuple or record."
                  </> textwrap "Consider ascribing an explicit type that does not reference "
                <> dquotes (prettyName d)
                <> "."
        | otherwise = verifyParams forbidden' ps
        where
          forbidden' =
            case patternParam p of
              (Named v, _, _) -> delete v forbidden
              _ -> forbidden
      verifyParams _ [] = 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 [] ret = RetType [] ret
  injectExt ext ret = RetType ext_here $ deeper ret
    where
      (immediate, _) = dimUses ret
      (ext_here, ext_there) = partition (`S.member` immediate) ext
      deeper :: TypeBase Size u -> TypeBase Size u
      deeper (Scalar (Prim t)) = Scalar $ Prim t
      deeper (Scalar (Record fs)) = Scalar $ Record $ M.map deeper fs
      deeper (Scalar (Sum cs)) = Scalar $ Sum $ M.map (map deeper) cs
      deeper (Scalar (Arrow als p d1 t1 (RetType t2_ext t2))) =
        Scalar $ Arrow als p d1 t1 $ injectExt (nubOrd (ext_there <> t2_ext)) t2
      deeper (Scalar (TypeVar u tn targs)) =
        Scalar $ TypeVar u tn $ map deeperArg targs
      deeper t@Array {} = t
  
      deeperArg (TypeArgType t) = TypeArgType $ deeper t
      deeperArg (TypeArgDim d) = TypeArgDim 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 defname defloc tparams paramts ret substs = do
    (more_tparams, retext) <-
      partitionEithers . catMaybes
        <$> mapM closeOver (M.toList $ M.map snd to_close_over)
    let mkExt v =
          case M.lookup v substs of
            Just (_, UnknownSize {}) -> Just v
            _ -> Nothing
    pure
      ( tparams ++ more_tparams,
        injectExt (nubOrd $ retext ++ mapMaybe mkExt (S.toList $ fvVars $ freeInType ret)) ret
      )
    where
      -- Diet does not matter here.
      t = foldFunType (map (toParam Observe) paramts) $ RetType [] ret
      to_close_over = M.filterWithKey (\k _ -> k `S.member` visible) substs
      visible = typeVars t <> fvVars (freeInType t)
  
      (produced_sizes, param_sizes) = dimUses t
  
      -- Avoid duplicate type parameters.
      closeOver (k, _)
        | k `elem` map typeParamName tparams =
            pure Nothing
      closeOver (k, NoConstraint l usage) =
        pure $ Just $ Left $ TypeParamType l k $ srclocOf usage
      closeOver (k, ParamType l loc) =
        pure $ Just $ Left $ TypeParamType l k $ srclocOf loc
      closeOver (k, Size Nothing usage) =
        pure $ Just $ Left $ TypeParamDim k $ srclocOf usage
      closeOver (k, UnknownSize _ _)
        | k `S.member` param_sizes,
          k `S.notMember` produced_sizes = do
            notes <- dimNotes defloc $ sizeFromName (qualName k) mempty
            typeError defloc notes . withIndexLink "unknown-param-def" $
              "Unknown size"
                <+> dquotes (prettyName k)
                <+> "in parameter of"
                <+> dquotes (prettyName defname)
                <> ", which is inferred as:"
                  </> indent 2 (pretty t)
        | k `S.member` produced_sizes =
            pure $ Just $ Right k
      closeOver (_, _) =
        pure Nothing
  
  letGeneralise ::
    Name ->
    SrcLoc ->
    [TypeParam] ->
    [Pat ParamType] ->
    ResType ->
    TermTypeM ([TypeParam], [Pat ParamType], ResRetType)
  letGeneralise defname defloc tparams params restype =
    onFailure (CheckingLetGeneralise defname) $ do
      now_substs <- 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 = overloadedTypeVars now_substs
  
      cur_lvl <- curLevel
      let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= (cur_lvl - length params)
          new_substs = M.filterWithKey candidate now_substs
  
      (tparams', RetType ret_dims restype') <-
        closeOverTypes
          defname
          defloc
          tparams
          (map patternStructType params)
          restype
          new_substs
  
      restype'' <- updateTypes restype'
  
      let used_sizes =
            freeInType restype'' <> foldMap (freeInType . patternType) params
      case filter ((`S.notMember` fvVars used_sizes) . typeParamName) $
        filter isSizeParam tparams' of
        [] -> pure ()
        tp : _ -> unusedSize $ SizeBinder (typeParamName tp) (srclocOf tp)
  
      -- We keep those type variables that were not closed over by
      -- let-generalisation.
      modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams'
  
      pure (tparams', params, RetType ret_dims restype'')
  
  checkFunBody ::
    [Pat ParamType] ->
    ExpBase NoInfo VName ->
    Maybe ResType ->
    SrcLoc ->
    TermTypeM Exp
  checkFunBody params body maybe_rettype loc = do
    body' <- checkExp body
  
    -- Unify body return type with return annotation, if one exists.
    case maybe_rettype of
      Just rettype -> do
        body_t <- expTypeFully body'
        -- We need to turn any sizes provided by "hidden" parameter
        -- names into existential sizes instead.
        let hidden = hiddenParamNames params
        (body_t', _) <-
          unscopeType
            loc
            (filter (`elem` hidden) $ foldMap patNames params)
            body_t
        case find (`elem` hidden) $ fvVars $ freeInType rettype of
          Just v ->
            typeError loc mempty $
              "The return type annotation"
                </> indent 2 (align (pretty rettype))
                </> "refers to the name"
                <+> dquotes (prettyName v)
                <+> "which is bound to an inner component of a function parameter."
          Nothing -> do
            let usage = mkUsage body "return type annotation"
            onFailure (CheckingReturn rettype body_t') $
              unify usage (toStruct rettype) body_t'
      Nothing -> pure ()
  
    pure body'
  
  arrayOfM ::
    SrcLoc ->
    StructType ->
    Shape Size ->
    TermTypeM StructType
  arrayOfM loc t shape = do
    arrayElemType (mkUsage loc "use as array element") "type used in array" t
    pure $ arrayOf shape t
  