content: -- | Main monad in which the type checker runs, as well as ancillary
  -- data definitions.
  module Language.Futhark.TypeChecker.Monad
    ( TypeM,
      runTypeM,
      askEnv,
      askImportName,
      atTopLevel,
      enteringModule,
      bindSpaced,
      bindSpaced1,
      bindIdents,
      qualifyTypeVars,
      lookupMTy,
      lookupImport,
      lookupMod,
      localEnv,
      TypeError (..),
      prettyTypeError,
      prettyTypeErrorNoLoc,
      withIndexLink,
      unappliedFunctor,
      unknownVariable,
      underscoreUse,
      Notes,
      aNote,
      MonadTypeChecker (..),
      TypeState (stateNameSource),
      usedName,
      checkName,
      checkAttr,
      checkQualName,
      checkValName,
      badOnLeft,
      isKnownType,
      module Language.Futhark.Warnings,
      Env (..),
      TySet,
      FunModType (..),
      ImportTable,
      NameMap,
      BoundV (..),
      Mod (..),
      TypeBinding (..),
      MTy (..),
      anySignedType,
      anyUnsignedType,
      anyIntType,
      anyFloatType,
      anyNumberType,
      anyPrimType,
      Namespace (..),
      intrinsicsNameMap,
      topLevelNameMap,
      mkTypeVarName,
    )
  where
  
  import Control.Monad
  import Control.Monad.Except
  import Control.Monad.Identity
  import Control.Monad.Reader
  import Control.Monad.State.Strict
  import Data.Either
  import Data.List (find)
  import Data.Map.Strict qualified as M
  import Data.Maybe
  import Data.Set qualified as S
  import Data.Text qualified as T
  import Data.Version qualified as Version
  import Futhark.FreshNames hiding (newName)
  import Futhark.FreshNames qualified
  import Futhark.Util.Pretty hiding (space)
  import Language.Futhark
  import Language.Futhark.Semantic
  import Language.Futhark.Traversals
  import Language.Futhark.Warnings
  import Paths_futhark qualified
  import Prelude hiding (mapM, mod)
  
  newtype Note = Note (Doc ())
  
  -- | A collection of extra information regarding a type error.
  newtype Notes = Notes [Note]
    deriving (Semigroup, Monoid)
  
  instance Pretty Note where
    pretty (Note msg) = unAnnotate $ "Note:" <+> align msg
  
  instance Pretty Notes where
    pretty (Notes notes) = unAnnotate $ foldMap (((line <> line) <>) . pretty) notes
  
  -- | A single note.
  aNote :: Doc () -> Notes
  aNote = Notes . pure . Note
  
  -- | Information about an error during type checking.
  data TypeError = TypeError Loc Notes (Doc ())
  
  -- | Prettyprint type error.
  prettyTypeError :: TypeError -> Doc AnsiStyle
  prettyTypeError (TypeError loc notes msg) =
    annotate
      (bold <> color Red)
      ("Error at " <> pretty (locText (srclocOf loc)) <> ":")
      </> prettyTypeErrorNoLoc (TypeError loc notes msg)
  
  -- | Prettyprint type error, without location information.  This can
  -- be used for cases where the location is printed in some other way.
  prettyTypeErrorNoLoc :: TypeError -> Doc AnsiStyle
  prettyTypeErrorNoLoc (TypeError _ notes msg) =
    unAnnotate msg <> pretty notes <> hardline
  
  errorIndexUrl :: Doc a
  errorIndexUrl = version_url <> "error-index.html"
    where
      version = Paths_futhark.version
      base_url = "https://futhark.readthedocs.io/en/"
      version_url
        | last (Version.versionBranch version) == 0 = base_url <> "latest/"
        | otherwise = base_url <> "v" <> pretty (Version.showVersion version) <> "/"
  
  -- | Attach a reference to documentation explaining the error in more detail.
  withIndexLink :: Doc a -> Doc a -> Doc a
  withIndexLink href msg =
    stack
      [ msg,
        "\nFor more information, see:",
        indent 2 (errorIndexUrl <> "#" <> href)
      ]
  
  -- | An unexpected functor appeared!
  unappliedFunctor :: (MonadTypeChecker m) => SrcLoc -> m a
  unappliedFunctor loc =
    typeError loc mempty "Cannot have parametric module here."
  
  -- | An unknown variable was referenced.
  unknownVariable ::
    (MonadTypeChecker m) =>
    Namespace ->
    QualName Name ->
    SrcLoc ->
    m a
  unknownVariable space name loc =
    typeError loc mempty $
      "Unknown" <+> pretty space <+> dquotes (pretty name)
  
  -- | A name prefixed with an underscore was used.
  underscoreUse ::
    (MonadTypeChecker m) =>
    SrcLoc ->
    QualName Name ->
    m a
  underscoreUse loc name =
    typeError loc mempty $
      "Use of"
        <+> dquotes (pretty name)
        <> ": variables prefixed with underscore may not be accessed."
  
  -- | A mapping from import import names to 'Env's.  This is used to
  -- resolve @import@ declarations.
  type ImportTable = M.Map ImportName Env
  
  data Context = Context
    { contextEnv :: Env,
      contextImportTable :: ImportTable,
      contextImportName :: ImportName,
      -- | Currently type-checking at the top level?  If false, we are
      -- inside a module.
      contextAtTopLevel :: Bool
    }
  
  data TypeState = TypeState
    { stateNameSource :: VNameSource,
      stateWarnings :: Warnings,
      -- | Which names have been used.
      stateUsed :: S.Set VName,
      stateCounter :: Int
    }
  
  -- | The type checker runs in this monad.
  newtype TypeM a
    = TypeM
        ( ReaderT
            Context
            (StateT TypeState (Except (Warnings, TypeError)))
            a
        )
    deriving
      ( Monad,
        Functor,
        Applicative,
        MonadReader Context,
        MonadState TypeState
      )
  
  instance MonadError TypeError TypeM where
    throwError e = TypeM $ do
      ws <- gets stateWarnings
      throwError (ws, e)
  
    catchError (TypeM m) f =
      TypeM $ m `catchError` f'
      where
        f' (_, e) =
          let TypeM m' = f e
           in m'
  
  -- | Run a 'TypeM' computation.
  runTypeM ::
    Env ->
    ImportTable ->
    ImportName ->
    VNameSource ->
    TypeM a ->
    (Warnings, Either TypeError (a, VNameSource))
  runTypeM env imports fpath src (TypeM m) = do
    let ctx = Context env imports fpath True
        s = TypeState src mempty mempty 0
    case runExcept $ runStateT (runReaderT m ctx) s of
      Left (ws, e) -> (ws, Left e)
      Right (x, s') -> (stateWarnings s', Right (x, stateNameSource s'))
  
  -- | Retrieve the current 'Env'.
  askEnv :: TypeM Env
  askEnv = asks contextEnv
  
  -- | The name of the current file/import.
  askImportName :: TypeM ImportName
  askImportName = asks contextImportName
  
  -- | Are we type-checking at the top level, or are we inside a nested
  -- module?
  atTopLevel :: TypeM Bool
  atTopLevel = asks contextAtTopLevel
  
  -- | We are now going to type-check the body of a module.
  enteringModule :: TypeM a -> TypeM a
  enteringModule = local $ \ctx -> ctx {contextAtTopLevel = False}
  
  -- | Look up a module type.
  lookupMTy :: SrcLoc -> QualName Name -> TypeM (QualName VName, MTy)
  lookupMTy loc qn = do
    (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Signature qn loc
    (qn',) <$> maybe explode pure (M.lookup name $ envModTypeTable scope)
    where
      explode = unknownVariable Signature qn loc
  
  -- | Look up an import.
  lookupImport :: SrcLoc -> FilePath -> TypeM (ImportName, Env)
  lookupImport loc file = do
    imports <- asks contextImportTable
    my_path <- asks contextImportName
    let canonical_import = mkImportFrom my_path file
    case M.lookup canonical_import imports of
      Nothing ->
        typeError loc mempty $
          "Unknown import"
            <+> dquotes (pretty (includeToText canonical_import))
            </> "Known:"
            <+> commasep (map (pretty . includeToText) (M.keys imports))
      Just scope -> pure (canonical_import, scope)
  
  -- | Evaluate a 'TypeM' computation within an extended (/not/
  -- replaced) environment.
  localEnv :: Env -> TypeM a -> TypeM a
  localEnv env = local $ \ctx ->
    let env' = env <> contextEnv ctx
     in ctx {contextEnv = env'}
  
  incCounter :: TypeM Int
  incCounter = do
    s <- get
    put s {stateCounter = stateCounter s + 1}
    pure $ stateCounter s
  
  bindNameMap :: NameMap -> TypeM a -> TypeM a
  bindNameMap m = local $ \ctx ->
    let env = contextEnv ctx
     in ctx {contextEnv = env {envNameMap = m <> envNameMap env}}
  
  -- | Monads that support type checking.  The reason we have this
  -- internal interface is because we use distinct monads for checking
  -- expressions and declarations.
  class (Monad m) => MonadTypeChecker m where
    warn :: (Located loc) => loc -> Doc () -> m ()
    warnings :: Warnings -> m ()
  
    newName :: VName -> m VName
    newID :: Name -> m VName
    newID s = newName $ VName s 0
    newTypeName :: Name -> m VName
  
    bindVal :: VName -> BoundV -> m a -> m a
  
    lookupType :: QualName VName -> m ([TypeParam], StructRetType, Liftedness)
  
    typeError :: (Located loc) => loc -> Notes -> Doc () -> m a
  
  warnIfUnused :: (Namespace, VName, SrcLoc) -> TypeM ()
  warnIfUnused (ns, name, loc) = do
    used <- gets stateUsed
    unless (name `S.member` used || "_" `T.isPrefixOf` nameToText (baseName name)) $
      warn loc $
        "Unused" <+> pretty ns <+> dquotes (prettyName name) <> "."
  
  -- | Map source-level names to fresh unique internal names, and
  -- evaluate a type checker context with the mapping active.
  bindSpaced :: [(Namespace, Name, SrcLoc)] -> ([VName] -> TypeM a) -> TypeM a
  bindSpaced names body = do
    names' <- mapM (\(_, v, _) -> newID v) names
    let mapping = M.fromList $ zip (map (\(ns, v, _) -> (ns, v)) names) $ map qualName names'
    bindNameMap mapping (body names')
      <* mapM_ warnIfUnused [(ns, v, loc) | ((ns, _, loc), v) <- zip names names']
  
  -- | Map single source-level name to fresh unique internal names, and
  -- evaluate a type checker context with the mapping active.
  bindSpaced1 :: Namespace -> Name -> SrcLoc -> (VName -> TypeM a) -> TypeM a
  bindSpaced1 ns name loc body = do
    name' <- newID name
    let mapping = M.singleton (ns, name) $ qualName name'
    bindNameMap mapping (body name') <* warnIfUnused (ns, name', loc)
  
  -- | Bind these identifiers in the name map and also check whether
  -- they have been used.
  bindIdents :: [IdentBase NoInfo VName t] -> TypeM a -> TypeM a
  bindIdents idents body = do
    let mapping =
          M.fromList $
            zip
              (map ((Term,) . (baseName . identName)) idents)
              (map (qualName . identName) idents)
    bindNameMap mapping body <* mapM_ warnIfUnused [(Term, v, loc) | Ident v _ loc <- idents]
  
  -- | Indicate that this name has been used. This is usually done
  -- implicitly by other operations, but sometimes we want to make a
  -- "fake" use to avoid things like top level functions being
  -- considered unused.
  usedName :: VName -> TypeM ()
  usedName name = modify $ \s -> s {stateUsed = S.insert name $ stateUsed s}
  
  instance MonadTypeChecker TypeM where
    warnings ws =
      modify $ \s -> s {stateWarnings = stateWarnings s <> ws}
  
    warn loc problem =
      warnings $ singleWarning (locOf loc) problem
  
    newName v = do
      s <- get
      let (v', src') = Futhark.FreshNames.newName (stateNameSource s) v
      put $ s {stateNameSource = src'}
      pure v'
  
    newTypeName name = do
      i <- incCounter
      newID $ mkTypeVarName name i
  
    bindVal v t = local $ \ctx ->
      ctx
        { contextEnv =
            (contextEnv ctx)
              { envVtable = M.insert v t $ envVtable $ contextEnv ctx
              }
        }
  
    lookupType qn = do
      outer_env <- askEnv
      scope <- lookupQualNameEnv qn
      case M.lookup (qualLeaf qn) $ envTypeTable scope of
        Nothing -> error $ "lookupType: " <> show qn
        Just (TypeAbbr l ps (RetType dims def)) ->
          pure (ps, RetType dims $ qualifyTypeVars outer_env mempty (qualQuals qn) def, l)
  
    typeError loc notes s = throwError $ TypeError (locOf loc) notes s
  
  lookupQualNameEnv :: QualName VName -> TypeM Env
  lookupQualNameEnv qn@(QualName quals _) = do
    env <- askEnv
    descend env quals
    where
      descend scope [] = pure scope
      descend scope (q : qs)
        | Just (ModEnv q_scope) <- M.lookup q $ envModTable scope =
            descend q_scope qs
        | otherwise =
            error $ "lookupQualNameEnv: " ++ show qn
  
  checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
  checkQualNameWithEnv space qn@(QualName quals name) loc = do
    env <- askEnv
    descend env quals
    where
      descend scope []
        | Just name' <- M.lookup (space, name) $ envNameMap scope = do
            usedName $ qualLeaf name'
            pure (scope, name')
        | otherwise =
            unknownVariable space qn loc
      descend scope (q : qs)
        | Just (QualName _ q') <- M.lookup (Term, q) $ envNameMap scope,
          Just res <- M.lookup q' $ envModTable scope = do
            usedName q'
            case res of
              ModEnv q_scope -> do
                (scope', QualName qs' name') <- descend q_scope qs
                pure (scope', QualName (q' : qs') name')
              ModFun {} -> unappliedFunctor loc
        | otherwise =
            unknownVariable space qn loc
  
  -- | Elaborate the given qualified name in the given namespace at the
  -- given location, producing the corresponding unique 'QualName'.
  -- Fails if the name is a module.
  checkValName :: QualName Name -> SrcLoc -> TypeM (QualName VName)
  checkValName name loc = do
    (env, name') <- checkQualNameWithEnv Term name loc
    case M.lookup (qualLeaf name') $ envModTable env of
      Just _ -> unknownVariable Term name loc
      Nothing -> pure name'
  
  -- | Elaborate the given qualified name in the given namespace at the
  -- given location, producing the corresponding unique 'QualName'.
  checkQualName :: Namespace -> QualName Name -> SrcLoc -> TypeM (QualName VName)
  checkQualName space name loc = snd <$> checkQualNameWithEnv space name loc
  
  -- | Elaborate the given name in the given namespace at the given
  -- location, producing the corresponding unique 'VName'.
  checkName :: Namespace -> Name -> SrcLoc -> TypeM VName
  checkName space name loc = qualLeaf <$> checkQualName space (qualName name) loc
  
  -- | Does a type with this name already exist? This is used for
  -- warnings, so it is OK it's a little unprincipled.
  isKnownType :: QualName VName -> TypeM Bool
  isKnownType qn = do
    env <- askEnv
    descend env (qualQuals qn) (qualLeaf qn)
    where
      descend env [] v
        | Just v' <- M.lookup (Type, baseName v) $ envNameMap env =
            pure $ M.member (qualLeaf v') $ envTypeTable env
      descend env (q : qs) v
        | Just q' <- M.lookup (Term, baseName q) $ envNameMap env,
          Just (ModEnv env') <- M.lookup (qualLeaf q') $ envModTable env =
            descend env' qs v
      descend _ _ _ = pure False
  
  lookupMod :: SrcLoc -> QualName Name -> TypeM (QualName VName, Mod)
  lookupMod loc qn = do
    (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Term qn loc
    case M.lookup name $ envModTable scope of
      Nothing -> unknownVariable Term qn loc
      Just m -> pure (qn', m)
  
  -- | Try to prepend qualifiers to the type names such that they
  -- represent how to access the type in some scope.
  qualifyTypeVars ::
    Env ->
    [VName] ->
    [VName] ->
    TypeBase Size as ->
    TypeBase Size as
  qualifyTypeVars outer_env orig_except ref_qs = onType (S.fromList orig_except)
    where
      onType ::
        S.Set VName ->
        TypeBase Size as ->
        TypeBase Size as
      onType except (Array u shape et) =
        Array u (fmap (onDim except) shape) (onScalar except et)
      onType except (Scalar t) =
        Scalar $ onScalar except t
  
      onScalar _ (Prim t) = Prim t
      onScalar except (TypeVar u qn targs) =
        TypeVar u (qual except qn) (map (onTypeArg except) targs)
      onScalar except (Record m) =
        Record $ M.map (onType except) m
      onScalar except (Sum m) =
        Sum $ M.map (map $ onType except) m
      onScalar except (Arrow as p d t1 (RetType dims t2)) =
        Arrow as p d (onType except' t1) $ RetType dims (onType except' t2)
        where
          except' = case p of
            Named p' -> S.insert p' except
            Unnamed -> except
  
      onTypeArg except (TypeArgDim d) =
        TypeArgDim $ onDim except d
      onTypeArg except (TypeArgType t) =
        TypeArgType $ onType except t
  
      onDim except e = runIdentity $ onDimM except e
      onDimM except (Var qn typ loc) = pure $ Var (qual except qn) typ loc
      onDimM except e = astMap (identityMapper {mapOnExp = onDimM except}) e
  
      qual except (QualName orig_qs name)
        | name `elem` except || reachable orig_qs name outer_env =
            QualName orig_qs name
        | otherwise =
            prependAsNecessary [] ref_qs $ QualName orig_qs name
  
      prependAsNecessary qs rem_qs (QualName orig_qs name)
        | reachable (qs ++ orig_qs) name outer_env = QualName (qs ++ orig_qs) name
        | otherwise = case rem_qs of
            q : rem_qs' -> prependAsNecessary (qs ++ [q]) rem_qs' (QualName orig_qs name)
            [] -> QualName orig_qs name
  
      reachable [] name env =
        name `M.member` envVtable env
          || isJust (find matches $ M.elems (envTypeTable env))
        where
          matches (TypeAbbr _ _ (RetType _ (Scalar (TypeVar _ (QualName x_qs name') _)))) =
            null x_qs && name == name'
          matches _ = False
      reachable (q : qs') name env
        | Just (ModEnv env') <- M.lookup q $ envModTable env =
            reachable qs' name env'
        | otherwise = False
  
  -- | Turn a 'Left' 'TypeError' into an actual error.
  badOnLeft :: Either TypeError a -> TypeM a
  badOnLeft = either throwError pure
  
  -- | All signed integer types.
  anySignedType :: [PrimType]
  anySignedType = map Signed [minBound .. maxBound]
  
  -- | All unsigned integer types.
  anyUnsignedType :: [PrimType]
  anyUnsignedType = map Unsigned [minBound .. maxBound]
  
  -- | All integer types.
  anyIntType :: [PrimType]
  anyIntType = anySignedType ++ anyUnsignedType
  
  -- | All floating-point types.
  anyFloatType :: [PrimType]
  anyFloatType = map FloatType [minBound .. maxBound]
  
  -- | All number types.
  anyNumberType :: [PrimType]
  anyNumberType = anyIntType ++ anyFloatType
  
  -- | All primitive types.
  anyPrimType :: [PrimType]
  anyPrimType = Bool : anyIntType ++ anyFloatType
  
  --- Name handling
  
  -- | The 'NameMap' corresponding to the intrinsics module.
  intrinsicsNameMap :: NameMap
  intrinsicsNameMap = M.fromList $ map mapping $ M.toList intrinsics
    where
      mapping (v, IntrinsicType {}) = ((Type, baseName v), QualName [] v)
      mapping (v, _) = ((Term, baseName v), QualName [] v)
  
  -- | The names that are available in the initial environment.
  topLevelNameMap :: NameMap
  topLevelNameMap = M.filterWithKey (\k _ -> available k) intrinsicsNameMap
    where
      available :: (Namespace, Name) -> Bool
      available (Type, _) = True
      available (Term, v) = v `S.member` (type_names <> binop_names <> fun_names)
        where
          type_names = S.fromList $ map (nameFromText . prettyText) anyPrimType
          binop_names =
            S.fromList $
              map
                (nameFromText . prettyText)
                [minBound .. (maxBound :: BinOp)]
          fun_names = S.fromList [nameFromString "shape"]
      available _ = False
  
  -- | Construct the name of a new type variable given a base
  -- description and a tag number (note that this is distinct from
  -- actually constructing a VName; the tag here is intended for human
  -- consumption but the machine does not care).
  mkTypeVarName :: Name -> Int -> Name
  mkTypeVarName desc i =
    desc <> nameFromString (mapMaybe subscript (show i))
    where
      subscript = flip lookup $ zip "0123456789" "₀₁₂₃₄₅₆₇₈₉"
  
  -- | Type-check an attribute.
  checkAttr :: (MonadTypeChecker m) => AttrInfo VName -> m (AttrInfo VName)
  checkAttr (AttrComp f attrs loc) =
    AttrComp f <$> mapM checkAttr attrs <*> pure loc
  checkAttr (AttrAtom (AtomName v) loc) =
    pure $ AttrAtom (AtomName v) loc
  checkAttr (AttrAtom (AtomInt x) loc) =
    pure $ AttrAtom (AtomInt x) loc
  