TypeChecker.hs

-- | The type checker checks whether the program is type-consistent
-- and adds type annotations and various other elaborations.  The
-- program does not need to have any particular properties for the
-- type checker to function; in particular it does not need unique
-- names.
module Language.Futhark.TypeChecker
  ( checkProg,
    checkExp,
    checkDec,
    checkModExp,
    Notes,
    TypeError (..),
    prettyTypeError,
    prettyTypeErrorNoLoc,
    Warnings,
    initialEnv,
    envWithImports,
  )
where

import Control.Monad
import Data.Bifunctor
import Data.Either
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Data.Set qualified as S
import Futhark.FreshNames hiding (newName)
import Futhark.Util.Pretty hiding (space)
import Language.Futhark
import Language.Futhark.Semantic
import Language.Futhark.TypeChecker.Modules
import Language.Futhark.TypeChecker.Monad
import Language.Futhark.TypeChecker.Names
import Language.Futhark.TypeChecker.Terms
import Language.Futhark.TypeChecker.Types
import Prelude hiding (abs, mod)

--- The main checker

-- | Type check a program containing no type information, yielding
-- either a type error or a program with complete type information.
-- Accepts a mapping from file names (excluding extension) to
-- previously type checked results.  The 'ImportName' is used to resolve
-- relative @import@s.
checkProg ::
  Imports ->
  VNameSource ->
  ImportName ->
  UncheckedProg ->
  (Warnings, Either TypeError (FileModule, VNameSource))
checkProg files src name prog =
  runTypeM initialEnv files' name src $ checkProgM prog
  where
    files' = M.map fileEnv $ M.fromList files

-- | Type check a single expression containing no type information,
-- yielding either a type error or the same expression annotated with
-- type information.  Also returns a list of type parameters, which
-- will be nonempty if the expression is polymorphic.  See also
-- 'checkProg'.
checkExp ::
  Imports ->
  VNameSource ->
  Env ->
  UncheckedExp ->
  (Warnings, Either TypeError ([TypeParam], Exp))
checkExp files src env e =
  second (fmap fst) $
    runTypeM
      env
      files'
      (mkInitialImport "")
      src
      (checkOneExp =<< resolveExp e)
  where
    files' = M.map fileEnv $ M.fromList files

-- | Type check a single declaration containing no type information,
-- yielding either a type error or the same declaration annotated with
-- type information along the Env produced by that declaration.  See
-- also 'checkProg'.
checkDec ::
  Imports ->
  VNameSource ->
  Env ->
  ImportName ->
  UncheckedDec ->
  (Warnings, Either TypeError (Env, Dec, VNameSource))
checkDec files src env name d =
  second (fmap massage) $
    runTypeM env files' name src $ do
      (_, env', d') <- checkOneDec d
      pure (env' <> env, d')
  where
    massage ((env', d'), src') =
      (env', d', src')
    files' = M.map fileEnv $ M.fromList files

-- | Type check a single module expression containing no type information,
-- yielding either a type error or the same expression annotated with
-- type information along the Env produced by that declaration.  See
-- also 'checkProg'.
checkModExp ::
  Imports ->
  VNameSource ->
  Env ->
  ModExpBase NoInfo Name ->
  (Warnings, Either TypeError (MTy, ModExpBase Info VName))
checkModExp files src env me =
  second (fmap fst) . runTypeM env files' (mkInitialImport "") src $ do
    (_abs, mty, me') <- checkOneModExp me
    pure (mty, me')
  where
    files' = M.map fileEnv $ M.fromList files

-- | An initial environment for the type checker, containing
-- intrinsics and such.
initialEnv :: Env
initialEnv =
  intrinsicsModule
    { envModTable = initialModTable,
      envNameMap =
        M.insert
          (Term, nameFromString "intrinsics")
          (qualName intrinsics_v)
          topLevelNameMap
    }
  where
    initialTypeTable = M.fromList $ mapMaybe addIntrinsicT $ M.toList intrinsics
    initialModTable = M.singleton intrinsics_v (ModEnv intrinsicsModule)

    intrinsics_v = VName (nameFromString "intrinsics") 0

    intrinsicsModule = Env mempty initialTypeTable mempty mempty intrinsicsNameMap

    addIntrinsicT (name, IntrinsicType l ps t) =
      Just (name, TypeAbbr l ps $ RetType [] t)
    addIntrinsicT _ =
      Nothing

-- | Produce an environment, based on the one passed in, where all of
-- the provided imports have been @open@ened in order.  This could in principle
-- also be done with 'checkDec', but this is more precise.
envWithImports :: Imports -> Env -> Env
envWithImports imports env =
  mconcat (map (fileEnv . snd) (reverse imports)) <> env

checkProgM :: UncheckedProg -> TypeM FileModule
checkProgM (Prog doc decs) = do
  checkForDuplicateDecs decs
  (abs, env, decs', full_env) <- checkDecs decs
  pure (FileModule abs env (Prog doc decs') full_env)

dupDefinitionError ::
  (MonadTypeChecker m) =>
  Namespace ->
  Name ->
  SrcLoc ->
  SrcLoc ->
  m a
dupDefinitionError space name loc1 loc2 =
  typeError loc1 mempty $
    "Duplicate definition of"
      <+> pretty space
      <+> prettyName name
      <> "."
        </> "Previously defined at"
        <+> pretty (locStr loc2)
      <> "."

checkForDuplicateDecs :: [DecBase NoInfo Name] -> TypeM ()
checkForDuplicateDecs =
  foldM_ (flip f) mempty
  where
    check namespace name loc known =
      case M.lookup (namespace, name) known of
        Just loc' ->
          dupDefinitionError namespace name loc loc'
        _ -> pure $ M.insert (namespace, name) loc known

    f (ValDec vb) =
      check Term (valBindName vb) (srclocOf vb)
    f (TypeDec (TypeBind name _ _ _ _ _ loc)) =
      check Type name loc
    f (ModTypeDec (ModTypeBind name _ _ loc)) =
      check Signature name loc
    f (ModDec (ModBind name _ _ _ _ loc)) =
      check Term name loc
    f OpenDec {} = pure
    f LocalDec {} = pure
    f ImportDec {} = pure

bindingTypeParams :: [TypeParam] -> TypeM a -> TypeM a
bindingTypeParams tparams = localEnv env
  where
    env = mconcat $ map typeParamEnv tparams

    typeParamEnv (TypeParamDim v _) =
      mempty
        { envVtable =
            M.singleton v $ BoundV [] (Scalar $ Prim $ Signed Int64)
        }
    typeParamEnv (TypeParamType l v _) =
      mempty
        { envTypeTable =
            M.singleton v $
              TypeAbbr l [] . RetType [] . Scalar $
                TypeVar mempty (qualName v) []
        }

checkTypeDecl ::
  UncheckedTypeExp ->
  TypeM ([VName], TypeExp Exp VName, StructType, Liftedness)
checkTypeDecl te = do
  (te', svars, RetType dims st, l) <- checkTypeExp checkSizeExp =<< resolveTypeExp te
  pure (svars ++ dims, te', toStruct st, l)

-- In this function, after the recursion, we add the Env of the
-- current Spec *after* the one that is returned from the recursive
-- call.  This implements the behaviour that specs later in a module
-- type can override those earlier (it rarely matters, but it affects
-- the specific structure of substitutions in case some module type is
-- redundantly imported multiple times).
checkSpecs :: [SpecBase NoInfo Name] -> TypeM (TySet, Env, [SpecBase Info VName])
checkSpecs [] = pure (mempty, mempty, [])
checkSpecs (ValSpec name tparams vtype NoInfo doc loc : specs) = do
  (tparams', vtype', vtype_t) <-
    resolveTypeParams tparams $ \tparams' -> bindingTypeParams tparams' $ do
      (ext, vtype', vtype_t, _) <- checkTypeDecl vtype

      unless (null ext) $
        typeError loc mempty $
          "All function parameters must have non-anonymous sizes."
            </> "Hint: add size parameters to"
            <+> dquotes (pretty name)
            <> "."

      pure (tparams', vtype', vtype_t)

  bindSpaced1 Term name loc $ \name' -> do
    let valenv =
          mempty
            { envVtable = M.singleton name' $ BoundV tparams' vtype_t,
              envNameMap = M.singleton (Term, name) $ qualName name'
            }
    usedName name'
    (abstypes, env, specs') <- localEnv valenv $ checkSpecs specs
    pure
      ( abstypes,
        env <> valenv,
        ValSpec name' tparams' vtype' (Info vtype_t) doc loc : specs'
      )
checkSpecs (TypeAbbrSpec tdec : specs) = do
  (tenv, tdec') <- checkTypeBind tdec
  bindSpaced1 Type (typeAlias tdec) (srclocOf tdec) $ \name' -> do
    usedName name'
    (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs
    pure
      ( abstypes,
        env <> tenv,
        TypeAbbrSpec tdec' : specs'
      )
checkSpecs (TypeSpec l name ps doc loc : specs) = do
  ps' <- resolveTypeParams ps pure
  bindSpaced1 Type name loc $ \name' -> do
    usedName name'
    let tenv =
          mempty
            { envNameMap =
                M.singleton (Type, name) $ qualName name',
              envTypeTable =
                M.singleton name' $
                  TypeAbbr l ps' . RetType [] . Scalar $
                    TypeVar mempty (qualName name') $
                      map typeParamToArg ps'
            }
    (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs
    pure
      ( M.insert (qualName name') l abstypes,
        env <> tenv,
        TypeSpec l name' ps' doc loc : specs'
      )
checkSpecs (ModSpec name sig doc loc : specs) = do
  (_sig_abs, mty, sig') <- checkModTypeExp sig
  bindSpaced1 Term name loc $ \name' -> do
    usedName name'
    let senv =
          mempty
            { envNameMap = M.singleton (Term, name) $ qualName name',
              envModTable = M.singleton name' $ mtyMod mty
            }
    (abstypes, env, specs') <- localEnv senv $ checkSpecs specs
    pure
      ( M.mapKeys (qualify name') (mtyAbs mty) <> abstypes,
        env <> senv,
        ModSpec name' sig' doc loc : specs'
      )
checkSpecs (IncludeSpec e loc : specs) = do
  (e_abs, env_abs, e_env, e') <- checkModTypeExpToEnv e

  mapM_ warnIfShadowing $ M.keys env_abs

  (abstypes, env, specs') <- localEnv e_env $ checkSpecs specs
  pure
    ( e_abs <> env_abs <> abstypes,
      env <> e_env,
      IncludeSpec e' loc : specs'
    )
  where
    warnIfShadowing qn = do
      known <- isKnownType qn
      when known $ warnAbout qn
    warnAbout qn =
      warn loc $ "Inclusion shadows type" <+> dquotes (pretty qn) <> "."

checkModTypeExp :: ModTypeExpBase NoInfo Name -> TypeM (TySet, MTy, ModTypeExpBase Info VName)
checkModTypeExp (ModTypeParens e loc) = do
  (abs, mty, e') <- checkModTypeExp e
  pure (abs, mty, ModTypeParens e' loc)
checkModTypeExp (ModTypeVar name NoInfo loc) = do
  (name', mty) <- lookupMTy loc name
  (mty', substs) <- newNamesForMTy mty
  pure (mtyAbs mty', mty', ModTypeVar name' (Info substs) loc)
checkModTypeExp (ModTypeSpecs specs loc) = do
  checkForDuplicateSpecs specs
  (abstypes, env, specs') <- checkSpecs specs
  pure (abstypes, MTy abstypes $ ModEnv env, ModTypeSpecs specs' loc)
checkModTypeExp (ModTypeWith s (TypeRef tname ps te trloc) loc) = do
  (abs, s_abs, s_env, s') <- checkModTypeExpToEnv s
  resolveTypeParams ps $ \ps' -> do
    (ext, te', te_t, _) <- bindingTypeParams ps' $ checkTypeDecl te
    (tname', s_abs', s_env') <-
      refineEnv loc s_abs s_env tname ps' $ RetType ext te_t
    pure (abs, MTy s_abs' $ ModEnv s_env', ModTypeWith s' (TypeRef tname' ps' te' trloc) loc)
checkModTypeExp (ModTypeArrow maybe_pname e1 e2 loc) = do
  (e1_abs, MTy s_abs e1_mod, e1') <- checkModTypeExp e1
  (maybe_pname', (e2_abs, e2_mod, e2')) <-
    case maybe_pname of
      Just pname -> bindSpaced1 Term pname loc $ \pname' ->
        localEnv (mempty {envModTable = M.singleton pname' e1_mod}) $
          (Just pname',) <$> checkModTypeExp e2
      Nothing ->
        (Nothing,) <$> checkModTypeExp e2
  pure
    ( e1_abs <> e2_abs,
      MTy mempty $ ModFun $ FunModType s_abs e1_mod e2_mod,
      ModTypeArrow maybe_pname' e1' e2' loc
    )

checkModTypeExpToEnv ::
  ModTypeExpBase NoInfo Name ->
  TypeM (TySet, TySet, Env, ModTypeExpBase Info VName)
checkModTypeExpToEnv e = do
  (abs, MTy mod_abs mod, e') <- checkModTypeExp e
  case mod of
    ModEnv env -> pure (abs, mod_abs, env, e')
    ModFun {} -> unappliedFunctor $ srclocOf e

checkModTypeBind :: ModTypeBindBase NoInfo Name -> TypeM (TySet, Env, ModTypeBindBase Info VName)
checkModTypeBind (ModTypeBind name e doc loc) = do
  (abs, env, e') <- checkModTypeExp e
  bindSpaced1 Signature name loc $ \name' -> do
    usedName name'
    pure
      ( abs,
        mempty
          { envModTypeTable = M.singleton name' env,
            envNameMap = M.singleton (Signature, name) (qualName name')
          },
        ModTypeBind name' e' doc loc
      )

checkOneModExp ::
  ModExpBase NoInfo Name ->
  TypeM (TySet, MTy, ModExpBase Info VName)
checkOneModExp (ModParens e loc) = do
  (abs, mty, e') <- checkOneModExp e
  pure (abs, mty, ModParens e' loc)
checkOneModExp (ModDecs decs loc) = do
  checkForDuplicateDecs decs
  (abstypes, env, decs', _) <- checkDecs decs
  pure
    ( abstypes,
      MTy abstypes $ ModEnv env,
      ModDecs decs' loc
    )
checkOneModExp (ModVar v loc) = do
  (v', env) <- lookupMod loc v
  when
    ( baseName (qualLeaf v') == nameFromString "intrinsics"
        && baseTag (qualLeaf v') <= maxIntrinsicTag
    )
    $ typeError loc mempty "The 'intrinsics' module may not be used in module expressions."
  pure (mempty, MTy mempty env, ModVar v' loc)
checkOneModExp (ModImport name NoInfo loc) = do
  (name', env) <- lookupImport loc name
  pure
    ( mempty,
      MTy mempty $ ModEnv env,
      ModImport name (Info name') loc
    )
checkOneModExp (ModApply f e NoInfo NoInfo loc) = do
  (f_abs, f_mty, f') <- checkOneModExp f
  case mtyMod f_mty of
    ModFun functor -> do
      (e_abs, e_mty, e') <- checkOneModExp e
      (mty, psubsts, rsubsts) <- applyFunctor (locOf loc) functor e_mty
      pure
        ( mtyAbs mty <> f_abs <> e_abs,
          mty,
          ModApply f' e' (Info psubsts) (Info rsubsts) loc
        )
    _ ->
      typeError loc mempty "Cannot apply non-parametric module."
checkOneModExp (ModAscript me se NoInfo loc) = do
  (me_abs, me_mod, me') <- checkOneModExp me
  (se_abs, se_mty, se') <- checkModTypeExp se
  match_subst <- badOnLeft $ matchMTys me_mod se_mty (locOf loc)
  pure (se_abs <> me_abs, se_mty, ModAscript me' se' (Info match_subst) loc)
checkOneModExp (ModLambda param maybe_fsig_e body_e loc) =
  withModParam param $ \param' param_abs param_mod -> do
    (abs, maybe_fsig_e', body_e', mty) <-
      checkModBody (fst <$> maybe_fsig_e) body_e loc
    pure
      ( abs,
        MTy mempty $ ModFun $ FunModType param_abs param_mod mty,
        ModLambda param' maybe_fsig_e' body_e' loc
      )

checkOneModExpToEnv :: ModExpBase NoInfo Name -> TypeM (TySet, Env, ModExpBase Info VName)
checkOneModExpToEnv e = do
  (e_abs, MTy abs mod, e') <- checkOneModExp e
  case mod of
    ModEnv env -> pure (e_abs <> abs, env, e')
    ModFun {} -> unappliedFunctor $ srclocOf e

withModParam ::
  ModParamBase NoInfo Name ->
  (ModParamBase Info VName -> TySet -> Mod -> TypeM a) ->
  TypeM a
withModParam (ModParam pname psig_e NoInfo loc) m = do
  (_abs, MTy p_abs p_mod, psig_e') <- checkModTypeExp psig_e
  bindSpaced1 Term pname loc $ \pname' -> do
    let in_body_env = mempty {envModTable = M.singleton pname' p_mod}
    localEnv in_body_env $
      m (ModParam pname' psig_e' (Info $ map qualLeaf $ M.keys p_abs) loc) p_abs p_mod

withModParams ::
  [ModParamBase NoInfo Name] ->
  ([(ModParamBase Info VName, TySet, Mod)] -> TypeM a) ->
  TypeM a
withModParams [] m = m []
withModParams (p : ps) m =
  withModParam p $ \p' pabs pmod ->
    withModParams ps $ \ps' -> m $ (p', pabs, pmod) : ps'

checkModBody ::
  Maybe (ModTypeExpBase NoInfo Name) ->
  ModExpBase NoInfo Name ->
  SrcLoc ->
  TypeM
    ( TySet,
      Maybe (ModTypeExp, Info (M.Map VName VName)),
      ModExp,
      MTy
    )
checkModBody maybe_fsig_e body_e loc = enteringModule $ do
  (body_e_abs, body_mty, body_e') <- checkOneModExp body_e
  case maybe_fsig_e of
    Nothing ->
      pure
        ( mtyAbs body_mty <> body_e_abs,
          Nothing,
          body_e',
          body_mty
        )
    Just fsig_e -> do
      (fsig_abs, fsig_mty, fsig_e') <- checkModTypeExp fsig_e
      fsig_subst <- badOnLeft $ matchMTys body_mty fsig_mty (locOf loc)
      pure
        ( fsig_abs <> body_e_abs,
          Just (fsig_e', Info fsig_subst),
          body_e',
          fsig_mty
        )

checkModBind :: ModBindBase NoInfo Name -> TypeM (TySet, Env, ModBindBase Info VName)
checkModBind (ModBind name [] maybe_fsig_e e doc loc) = do
  (e_abs, maybe_fsig_e', e', mty) <- checkModBody (fst <$> maybe_fsig_e) e loc
  bindSpaced1 Term name loc $ \name' -> do
    usedName name'
    pure
      ( e_abs,
        mempty
          { envModTable = M.singleton name' $ mtyMod mty,
            envNameMap = M.singleton (Term, name) $ qualName name'
          },
        ModBind name' [] maybe_fsig_e' e' doc loc
      )
checkModBind (ModBind name (p : ps) maybe_fsig_e body_e doc loc) = do
  (abs, params', maybe_fsig_e', body_e', funsig) <-
    withModParam p $ \p' p_abs p_mod ->
      withModParams ps $ \params_stuff -> do
        let (ps', ps_abs, ps_mod) = unzip3 params_stuff
        (abs, maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc
        let addParam (x, y) mty' = MTy mempty $ ModFun $ FunModType x y mty'
        pure
          ( abs,
            p' : ps',
            maybe_fsig_e',
            body_e',
            FunModType p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod
          )
  bindSpaced1 Term name loc $ \name' -> do
    usedName name'
    pure
      ( abs,
        mempty
          { envModTable =
              M.singleton name' $ ModFun funsig,
            envNameMap =
              M.singleton (Term, name) $ qualName name'
          },
        ModBind name' params' maybe_fsig_e' body_e' doc loc
      )

checkForDuplicateSpecs :: [SpecBase NoInfo Name] -> TypeM ()
checkForDuplicateSpecs =
  foldM_ (flip f) mempty
  where
    check namespace name loc known =
      case M.lookup (namespace, name) known of
        Just loc' ->
          dupDefinitionError namespace name loc loc'
        _ -> pure $ M.insert (namespace, name) loc known

    f (ValSpec name _ _ _ _ loc) =
      check Term name loc
    f (TypeAbbrSpec (TypeBind name _ _ _ _ _ loc)) =
      check Type name loc
    f (TypeSpec _ name _ _ loc) =
      check Type name loc
    f (ModSpec name _ _ loc) =
      check Term name loc
    f IncludeSpec {} =
      pure

checkTypeBind ::
  TypeBindBase NoInfo Name ->
  TypeM (Env, TypeBindBase Info VName)
checkTypeBind (TypeBind name l tps te NoInfo doc loc) =
  resolveTypeParams tps $ \tps' -> do
    (te', svars, RetType dims t, l') <-
      bindingTypeParams tps' $ checkTypeExp checkSizeExp =<< resolveTypeExp te

    let (witnessed, _) = determineSizeWitnesses $ toStruct t
    case L.find (`S.notMember` witnessed) svars of
      Just _ ->
        typeError (locOf te) mempty . withIndexLink "anonymous-nonconstructive" $
          "Type abbreviation contains an anonymous size not used constructively as an array size."
      Nothing ->
        pure ()

    let elab_t = RetType (svars ++ dims) $ toStruct t

    let used_dims = fvVars $ freeInType t
    case filter ((`S.notMember` used_dims) . typeParamName) $
      filter isSizeParam tps' of
      [] -> pure ()
      tp : _ ->
        typeError loc mempty $
          "Size parameter" <+> dquotes (pretty tp) <+> "unused."

    case (l, l') of
      (_, Lifted)
        | l < Lifted ->
            typeError loc mempty $
              "Non-lifted type abbreviations may not contain functions."
                </> "Hint: consider using 'type^'."
      (_, SizeLifted)
        | l < SizeLifted ->
            typeError loc mempty $
              "Non-size-lifted type abbreviations may not contain size-lifted types."
                </> "Hint: consider using 'type~'."
      (Unlifted, _)
        | not $ null $ svars ++ dims ->
            typeError loc mempty $
              "Non-lifted type abbreviations may not use existential sizes in their definition."
                </> "Hint: use 'type~' or add size parameters to"
                <+> dquotes (prettyName name)
                <> "."
      _ -> pure ()

    bindSpaced1 Type name loc $ \name' -> do
      usedName name'
      pure
        ( mempty
            { envTypeTable =
                M.singleton name' $ TypeAbbr l tps' elab_t,
              envNameMap =
                M.singleton (Type, name) $ qualName name'
            },
          TypeBind name' l tps' te' (Info elab_t) doc loc
        )

entryPoint :: [Pat ParamType] -> Maybe (TypeExp Exp VName) -> ResRetType -> EntryPoint
entryPoint params orig_ret_te (RetType _ret orig_ret) =
  EntryPoint (map patternEntry params ++ more_params) rettype'
  where
    (more_params, rettype') = onRetType orig_ret_te $ toStruct orig_ret

    patternEntry (PatParens p _) =
      patternEntry p
    patternEntry (PatAscription p te _) =
      EntryParam (patternName p) $ EntryType (patternStructType p) (Just te)
    patternEntry p =
      EntryParam (patternName p) $ EntryType (patternStructType p) Nothing

    patternName (Id x _ _) = baseName x
    patternName (PatParens p _) = patternName p
    patternName _ = "_"

    pname (Named v) = baseName v
    pname Unnamed = "_"
    onRetType (Just (TEArrow p t1_te t2_te _)) (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =
      let (xs, y) = onRetType (Just t2_te) $ toStruct t2
       in (EntryParam (maybe "_" baseName p) (EntryType t1 (Just t1_te)) : xs, y)
    onRetType _ (Scalar (Arrow _ p _ t1 (RetType _ t2))) =
      let (xs, y) = onRetType Nothing $ toStruct t2
       in (EntryParam (pname p) (EntryType t1 Nothing) : xs, y)
    onRetType te t =
      ([], EntryType t te)

checkEntryPoint ::
  SrcLoc ->
  [TypeParam] ->
  [Pat ParamType] ->
  ResRetType ->
  TypeM ()
checkEntryPoint loc tparams params rettype
  | any isTypeParam tparams =
      typeError loc mempty $
        withIndexLink
          "polymorphic-entry"
          "Entry point functions may not be polymorphic."
  | not (all orderZero param_ts)
      || not (orderZero rettype') =
      typeError loc mempty $
        withIndexLink
          "higher-order-entry"
          "Entry point functions may not be higher-order."
  | sizes_only_in_ret <-
      S.fromList (map typeParamName tparams)
        `S.intersection` fvVars (freeInType rettype')
        `S.difference` foldMap (fvVars . freeInType) param_ts,
    not $ S.null sizes_only_in_ret =
      typeError loc mempty $
        withIndexLink
          "size-polymorphic-entry"
          "Entry point functions must not be size-polymorphic in their return type."
  | (constructive, _) <- foldMap (determineSizeWitnesses . toStruct) param_ts,
    Just p <- L.find (flip S.notMember constructive . typeParamName) tparams =
      typeError p mempty . withIndexLink "nonconstructive-entry" $
        "Entry point size parameter "
          <> pretty p
          <> " only used non-constructively."
  | otherwise =
      pure ()
  where
    (RetType _ rettype_t) = rettype
    (rettype_params, rettype') = unfoldFunType rettype_t
    param_ts = map patternType params ++ rettype_params

checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind)
checkValBind vb = do
  (ValBind entry fname fname_loc maybe_tdecl NoInfo tparams params body doc attrs loc) <-
    resolveValBind vb

  top_level <- atTopLevel
  when (not top_level && isJust entry) $
    typeError loc mempty $
      withIndexLink "nested-entry" "Entry points may not be declared inside modules."

  attrs' <- mapM checkAttr attrs

  (tparams', params', maybe_tdecl', rettype, body') <-
    checkFunDef (fname, maybe_tdecl, tparams, params, body, loc)

  let entry' = Info (entryPoint params' maybe_tdecl' rettype) <$ entry
  case entry' of
    Just _ -> checkEntryPoint loc tparams' params' rettype
    _ -> pure ()

  let vb' = ValBind entry' fname fname_loc maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc
  pure
    ( mempty
        { envVtable =
            M.singleton fname $ uncurry BoundV $ valBindTypeScheme vb',
          envNameMap =
            M.singleton (Term, baseName fname) $ qualName fname
        },
      vb'
    )

checkOneDec :: DecBase NoInfo Name -> TypeM (TySet, Env, DecBase Info VName)
checkOneDec (ModDec struct) = do
  (abs, modenv, struct') <- checkModBind struct
  pure (abs, modenv, ModDec struct')
checkOneDec (ModTypeDec sig) = do
  (abs, sigenv, sig') <- checkModTypeBind sig
  pure (abs, sigenv, ModTypeDec sig')
checkOneDec (TypeDec tdec) = do
  (tenv, tdec') <- checkTypeBind tdec
  pure (mempty, tenv, TypeDec tdec')
checkOneDec (OpenDec x loc) = do
  (x_abs, x_env, x') <- checkOneModExpToEnv x
  pure (x_abs, x_env, OpenDec x' loc)
checkOneDec (LocalDec d loc) = do
  (abstypes, env, d') <- checkOneDec d
  pure (abstypes, env, LocalDec d' loc)
checkOneDec (ImportDec name NoInfo loc) = do
  (name', env) <- lookupImport loc name
  when (isBuiltin name) $
    typeError loc mempty $
      pretty name <+> "may not be explicitly imported."
  pure (mempty, env, ImportDec name (Info name') loc)
checkOneDec (ValDec vb) = do
  (env, vb') <- checkValBind vb
  pure (mempty, env, ValDec vb')

checkDecs :: [DecBase NoInfo Name] -> TypeM (TySet, Env, [DecBase Info VName], Env)
checkDecs (d : ds) = do
  (d_abstypes, d_env, d') <- checkOneDec d
  (ds_abstypes, ds_env, ds', full_env) <- localEnv d_env $ checkDecs ds
  pure
    ( d_abstypes <> ds_abstypes,
      case d' of
        LocalDec {} -> ds_env
        ImportDec {} -> ds_env
        _ -> ds_env <> d_env,
      d' : ds',
      full_env
    )
checkDecs [] = do
  full_env <- askEnv
  pure (mempty, mempty, [], full_env)