-- | 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,
    qualifyTypeVars,
    lookupMTy,
    lookupImport,
    localEnv,
    TypeError (..),
    prettyTypeError,
    prettyTypeErrorNoLoc,
    withIndexLink,
    unappliedFunctor,
    unknownVariable,
    unknownType,
    underscoreUse,
    Notes,
    aNote,
    MonadTypeChecker (..),
    checkName,
    checkAttr,
    badOnLeft,
    module Language.Futhark.Warnings,
    Env (..),
    TySet,
    FunSig (..),
    ImportTable,
    NameMap,
    BoundV (..),
    Mod (..),
    TypeBinding (..),
    MTy (..),
    anySignedType,
    anyUnsignedType,
    anyIntType,
    anyFloatType,
    anyNumberType,
    anyPrimType,
    Namespace (..),
    intrinsicsNameMap,
    topLevelNameMap,
    mkTypeVarName,
  )
where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Either
import Data.List (find, isPrefixOf)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
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.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 (NonEmpty Notes -> Notes
Notes -> Notes -> Notes
forall b. Integral b => b -> Notes -> Notes
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> Notes -> Notes
$cstimes :: forall b. Integral b => b -> Notes -> Notes
sconcat :: NonEmpty Notes -> Notes
$csconcat :: NonEmpty Notes -> Notes
<> :: Notes -> Notes -> Notes
$c<> :: Notes -> Notes -> Notes
Semigroup, Semigroup Notes
Notes
[Notes] -> Notes
Notes -> Notes -> Notes
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [Notes] -> Notes
$cmconcat :: [Notes] -> Notes
mappend :: Notes -> Notes -> Notes
$cmappend :: Notes -> Notes -> Notes
mempty :: Notes
$cmempty :: Notes
Monoid)

instance Pretty Note where
  pretty :: forall ann. Note -> Doc ann
pretty (Note Doc ()
msg) = forall ann xxx. Doc ann -> Doc xxx
unAnnotate forall a b. (a -> b) -> a -> b
$ Doc ()
"Note:" forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
align Doc ()
msg

instance Pretty Notes where
  pretty :: forall ann. Notes -> Doc ann
pretty (Notes [Note]
notes) = forall ann xxx. Doc ann -> Doc xxx
unAnnotate forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (((forall ann. Doc ann
line forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann
line) <>) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a ann. Pretty a => a -> Doc ann
pretty) [Note]
notes

-- | A single note.
aNote :: Doc () -> Notes
aNote :: Doc () -> Notes
aNote = [Note] -> Notes
Notes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Note
Note

-- | Information about an error during type checking.
data TypeError = TypeError Loc Notes (Doc ())

-- | Prettyprint type error.
prettyTypeError :: TypeError -> Doc AnsiStyle
prettyTypeError :: TypeError -> Doc AnsiStyle
prettyTypeError (TypeError Loc
loc Notes
notes Doc ()
msg) =
  forall ann. ann -> Doc ann -> Doc ann
annotate
    (AnsiStyle
bold forall a. Semigroup a => a -> a -> a
<> Color -> AnsiStyle
color Color
Red)
    (Doc AnsiStyle
"Error at " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty (forall a. Located a => a -> Text
locText (forall a. Located a => a -> SrcLoc
srclocOf Loc
loc)) forall a. Semigroup a => a -> a -> a
<> Doc AnsiStyle
":")
    forall ann. Doc ann -> Doc ann -> Doc ann
</> TypeError -> Doc AnsiStyle
prettyTypeErrorNoLoc (Loc -> Notes -> Doc () -> TypeError
TypeError Loc
loc Notes
notes Doc ()
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 -> Doc AnsiStyle
prettyTypeErrorNoLoc (TypeError Loc
_ Notes
notes Doc ()
msg) =
  forall ann xxx. Doc ann -> Doc xxx
unAnnotate Doc ()
msg forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Notes
notes forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann
hardline

errorIndexUrl :: Doc a
errorIndexUrl :: forall ann. Doc ann
errorIndexUrl = Doc a
version_url forall a. Semigroup a => a -> a -> a
<> Doc a
"error-index.html"
  where
    version :: Version
version = Version
Paths_futhark.version
    base_url :: Doc a
base_url = Doc a
"https://futhark.readthedocs.io/en/"
    version_url :: Doc a
version_url
      | forall a. [a] -> a
last (Version -> [Int]
Version.versionBranch Version
version) forall a. Eq a => a -> a -> Bool
== Int
0 = Doc a
base_url forall a. Semigroup a => a -> a -> a
<> Doc a
"latest/"
      | Bool
otherwise = Doc a
base_url forall a. Semigroup a => a -> a -> a
<> Doc a
"v" forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty (Version -> String
Version.showVersion Version
version) forall a. Semigroup a => a -> a -> a
<> Doc a
"/"

-- | Attach a reference to documentation explaining the error in more detail.
withIndexLink :: Doc a -> Doc a -> Doc a
withIndexLink :: forall ann. Doc ann -> Doc ann -> Doc ann
withIndexLink Doc a
href Doc a
msg =
  forall a. [Doc a] -> Doc a
stack
    [ Doc a
msg,
      Doc a
"\nFor more information, see:",
      forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall ann. Doc ann
errorIndexUrl forall a. Semigroup a => a -> a -> a
<> Doc a
"#" forall a. Semigroup a => a -> a -> a
<> Doc a
href)
    ]

-- | An unexpected functor appeared!
unappliedFunctor :: MonadTypeChecker m => SrcLoc -> m a
unappliedFunctor :: forall (m :: * -> *) a. MonadTypeChecker m => SrcLoc -> m a
unappliedFunctor SrcLoc
loc =
  forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty Doc ()
"Cannot have parametric module here."

-- | An unknown variable was referenced.
unknownVariable ::
  MonadTypeChecker m =>
  Namespace ->
  QualName Name ->
  SrcLoc ->
  m a
unknownVariable :: forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
name SrcLoc
loc =
  forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
    Doc ()
"Unknown" forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a ann. Pretty a => a -> Doc ann
pretty Namespace
space forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName Name
name)

-- | An unknown type was referenced.
unknownType :: MonadTypeChecker m => SrcLoc -> QualName Name -> m a
unknownType :: forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
unknownType SrcLoc
loc QualName Name
name =
  forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ Doc ()
"Unknown type" forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a ann. Pretty a => a -> Doc ann
pretty QualName Name
name forall a. Semigroup a => a -> a -> a
<> Doc ()
"."

-- | A name prefixed with an underscore was used.
underscoreUse ::
  MonadTypeChecker m =>
  SrcLoc ->
  QualName Name ->
  m a
underscoreUse :: forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
underscoreUse SrcLoc
loc QualName Name
name =
  forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
    Doc ()
"Use of"
      forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName Name
name)
        forall a. Semigroup a => a -> a -> a
<> Doc ()
": variables prefixed with underscore may not be accessed."

-- | A mapping from import strings to 'Env's.  This is used to resolve
-- @import@ declarations.
type ImportTable = M.Map String Env

data Context = Context
  { Context -> Env
contextEnv :: Env,
    Context -> ImportTable
contextImportTable :: ImportTable,
    Context -> ImportName
contextImportName :: ImportName,
    -- | Currently type-checking at the top level?  If false, we are
    -- inside a module.
    Context -> Bool
contextAtTopLevel :: Bool
  }

data TypeState = TypeState
  { TypeState -> VNameSource
stateNameSource :: VNameSource,
    TypeState -> Warnings
stateWarnings :: Warnings,
    TypeState -> Int
stateCounter :: Int
  }

-- | The type checker runs in this monad.
newtype TypeM a
  = TypeM
      ( ReaderT
          Context
          (StateT TypeState (Except (Warnings, TypeError)))
          a
      )
  deriving
    ( Applicative TypeM
forall a. a -> TypeM a
forall a b. TypeM a -> TypeM b -> TypeM b
forall a b. TypeM a -> (a -> TypeM b) -> TypeM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> TypeM a
$creturn :: forall a. a -> TypeM a
>> :: forall a b. TypeM a -> TypeM b -> TypeM b
$c>> :: forall a b. TypeM a -> TypeM b -> TypeM b
>>= :: forall a b. TypeM a -> (a -> TypeM b) -> TypeM b
$c>>= :: forall a b. TypeM a -> (a -> TypeM b) -> TypeM b
Monad,
      forall a b. a -> TypeM b -> TypeM a
forall a b. (a -> b) -> TypeM a -> TypeM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> TypeM b -> TypeM a
$c<$ :: forall a b. a -> TypeM b -> TypeM a
fmap :: forall a b. (a -> b) -> TypeM a -> TypeM b
$cfmap :: forall a b. (a -> b) -> TypeM a -> TypeM b
Functor,
      Functor TypeM
forall a. a -> TypeM a
forall a b. TypeM a -> TypeM b -> TypeM a
forall a b. TypeM a -> TypeM b -> TypeM b
forall a b. TypeM (a -> b) -> TypeM a -> TypeM b
forall a b c. (a -> b -> c) -> TypeM a -> TypeM b -> TypeM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. TypeM a -> TypeM b -> TypeM a
$c<* :: forall a b. TypeM a -> TypeM b -> TypeM a
*> :: forall a b. TypeM a -> TypeM b -> TypeM b
$c*> :: forall a b. TypeM a -> TypeM b -> TypeM b
liftA2 :: forall a b c. (a -> b -> c) -> TypeM a -> TypeM b -> TypeM c
$cliftA2 :: forall a b c. (a -> b -> c) -> TypeM a -> TypeM b -> TypeM c
<*> :: forall a b. TypeM (a -> b) -> TypeM a -> TypeM b
$c<*> :: forall a b. TypeM (a -> b) -> TypeM a -> TypeM b
pure :: forall a. a -> TypeM a
$cpure :: forall a. a -> TypeM a
Applicative,
      MonadReader Context,
      MonadState TypeState
    )

instance MonadError TypeError TypeM where
  throwError :: forall a. TypeError -> TypeM a
throwError TypeError
e = forall a.
ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
TypeM forall a b. (a -> b) -> a -> b
$ do
    Warnings
ws <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TypeState -> Warnings
stateWarnings
    forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Warnings
ws, TypeError
e)

  catchError :: forall a. TypeM a -> (TypeError -> TypeM a) -> TypeM a
catchError (TypeM ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m) TypeError -> TypeM a
f =
    forall a.
ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
TypeM forall a b. (a -> b) -> a -> b
$ ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m forall e (m :: * -> *) a.
MonadError e m =>
m a -> (e -> m a) -> m a
`catchError` forall {a}.
(a, TypeError)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
f'
    where
      f' :: (a, TypeError)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
f' (a
_, TypeError
e) =
        let TypeM ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m' = TypeError -> TypeM a
f TypeError
e
         in ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m'

-- | Run a 'TypeM' computation.
runTypeM ::
  Env ->
  ImportTable ->
  ImportName ->
  VNameSource ->
  TypeM a ->
  (Warnings, Either TypeError (a, VNameSource))
runTypeM :: forall a.
Env
-> ImportTable
-> ImportName
-> VNameSource
-> TypeM a
-> (Warnings, Either TypeError (a, VNameSource))
runTypeM Env
env ImportTable
imports ImportName
fpath VNameSource
src (TypeM ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m) = do
  let ctx :: Context
ctx = Env -> ImportTable -> ImportName -> Bool -> Context
Context Env
env ImportTable
imports ImportName
fpath Bool
True
      s :: TypeState
s = VNameSource -> Warnings -> Int -> TypeState
TypeState VNameSource
src forall a. Monoid a => a
mempty Int
0
  case forall e a. Except e a -> Either e a
runExcept forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m Context
ctx) TypeState
s of
    Left (Warnings
ws, TypeError
e) -> (Warnings
ws, forall a b. a -> Either a b
Left TypeError
e)
    Right (a
x, TypeState VNameSource
src' Warnings
ws Int
_) -> (Warnings
ws, forall a b. b -> Either a b
Right (a
x, VNameSource
src'))

-- | Retrieve the current 'Env'.
askEnv :: TypeM Env
askEnv :: TypeM Env
askEnv = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> Env
contextEnv

-- | The name of the current file/import.
askImportName :: TypeM ImportName
askImportName :: TypeM ImportName
askImportName = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> ImportName
contextImportName

-- | Are we type-checking at the top level, or are we inside a nested
-- module?
atTopLevel :: TypeM Bool
atTopLevel :: TypeM Bool
atTopLevel = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> Bool
contextAtTopLevel

-- | We are now going to type-check the body of a module.
enteringModule :: TypeM a -> TypeM a
enteringModule :: forall a. TypeM a -> TypeM a
enteringModule = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Context
ctx -> Context
ctx {contextAtTopLevel :: Bool
contextAtTopLevel = Bool
False}

-- | Look up a module type.
lookupMTy :: SrcLoc -> QualName Name -> TypeM (QualName VName, MTy)
lookupMTy :: SrcLoc -> QualName Name -> TypeM (QualName VName, MTy)
lookupMTy SrcLoc
loc QualName Name
qn = do
  (Env
scope, qn' :: QualName VName
qn'@(QualName [VName]
_ VName
name)) <- Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv Namespace
Signature QualName Name
qn SrcLoc
loc
  (QualName VName
qn',) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall {a}. TypeM a
explode forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall a b. (a -> b) -> a -> b
$ Env -> Map VName MTy
envSigTable Env
scope)
  where
    explode :: TypeM a
explode = forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
Signature QualName Name
qn SrcLoc
loc

-- | Look up an import.
lookupImport :: SrcLoc -> FilePath -> TypeM (FilePath, Env)
lookupImport :: SrcLoc -> String -> TypeM (String, Env)
lookupImport SrcLoc
loc String
file = do
  ImportTable
imports <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> ImportTable
contextImportTable
  ImportName
my_path <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> ImportName
contextImportName
  let canonical_import :: String
canonical_import = ImportName -> String
includeToString forall a b. (a -> b) -> a -> b
$ ImportName -> String -> SrcLoc -> ImportName
mkImportFrom ImportName
my_path String
file SrcLoc
loc
  case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
canonical_import ImportTable
imports of
    Maybe Env
Nothing ->
      forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
        Doc ()
"Unknown import"
          forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty String
canonical_import)
          forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"Known:"
          forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a. [Doc a] -> Doc a
commasep (forall a b. (a -> b) -> [a] -> [b]
map forall a ann. Pretty a => a -> Doc ann
pretty (forall k a. Map k a -> [k]
M.keys ImportTable
imports))
    Just Env
scope -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
canonical_import, Env
scope)

-- | Evaluate a 'TypeM' computation within an extended (/not/
-- replaced) environment.
localEnv :: Env -> TypeM a -> TypeM a
localEnv :: forall a. Env -> TypeM a -> TypeM a
localEnv Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
  let env' :: Env
env' = Env
env forall a. Semigroup a => a -> a -> a
<> Context -> Env
contextEnv Context
ctx
   in Context
ctx {contextEnv :: Env
contextEnv = Env
env'}

incCounter :: TypeM Int
incCounter :: TypeM Int
incCounter = do
  TypeState
s <- forall s (m :: * -> *). MonadState s m => m s
get
  forall s (m :: * -> *). MonadState s m => s -> m ()
put TypeState
s {stateCounter :: Int
stateCounter = TypeState -> Int
stateCounter TypeState
s forall a. Num a => a -> a -> a
+ Int
1}
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ TypeState -> Int
stateCounter TypeState
s

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

  newName :: VName -> m VName
  newID :: Name -> m VName
  newTypeName :: Name -> m VName

  bindNameMap :: NameMap -> m a -> m a
  bindVal :: VName -> BoundV -> m a -> m a

  checkQualName :: Namespace -> QualName Name -> SrcLoc -> m (QualName VName)

  lookupType :: SrcLoc -> QualName Name -> m (QualName VName, [TypeParam], StructRetType, Liftedness)
  lookupMod :: SrcLoc -> QualName Name -> m (QualName VName, Mod)
  lookupVar :: SrcLoc -> QualName Name -> m (QualName VName, PatType)

  checkNamedSize :: SrcLoc -> QualName Name -> m (QualName VName)
  checkNamedSize SrcLoc
loc QualName Name
v = do
    (QualName VName
v', PatType
t) <- forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m (QualName VName, PatType)
lookupVar SrcLoc
loc QualName Name
v
    case PatType
t of
      Scalar (Prim (Signed IntType
Int64)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure QualName VName
v'
      PatType
_ ->
        forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
          Doc ()
"Sizes must have type i64, but"
            forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName Name
v)
            forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
"has type:"
            forall ann. Doc ann -> Doc ann -> Doc ann
</> forall a ann. Pretty a => a -> Doc ann
pretty PatType
t

  typeError :: Located loc => loc -> Notes -> Doc () -> m a

-- | Elaborate the given name in the given namespace at the given
-- location, producing the corresponding unique 'VName'.
checkName :: MonadTypeChecker m => Namespace -> Name -> SrcLoc -> m VName
checkName :: forall (m :: * -> *).
MonadTypeChecker m =>
Namespace -> Name -> SrcLoc -> m VName
checkName Namespace
space Name
name SrcLoc
loc = forall vn. QualName vn -> vn
qualLeaf forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m (QualName VName)
checkQualName Namespace
space (forall v. v -> QualName v
qualName Name
name) SrcLoc
loc

-- | Map source-level names do fresh unique internal names, and
-- evaluate a type checker context with the mapping active.
bindSpaced :: MonadTypeChecker m => [(Namespace, Name)] -> m a -> m a
bindSpaced :: forall (m :: * -> *) a.
MonadTypeChecker m =>
[(Namespace, Name)] -> m a -> m a
bindSpaced [(Namespace, Name)]
names m a
body = do
  [VName]
names' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Namespace, Name)]
names
  let mapping :: Map (Namespace, Name) (QualName VName)
mapping = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (forall a b. [a] -> [b] -> [(a, b)]
zip [(Namespace, Name)]
names forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall v. v -> QualName v
qualName [VName]
names')
  forall (m :: * -> *) a.
MonadTypeChecker m =>
Map (Namespace, Name) (QualName VName) -> m a -> m a
bindNameMap Map (Namespace, Name) (QualName VName)
mapping m a
body

instance MonadTypeChecker TypeM where
  warn :: forall loc. Located loc => loc -> Doc () -> TypeM ()
warn loc
loc Doc ()
problem =
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \TypeState
s ->
      TypeState
s
        { stateWarnings :: Warnings
stateWarnings = TypeState -> Warnings
stateWarnings TypeState
s forall a. Semigroup a => a -> a -> a
<> SrcLoc -> Doc () -> Warnings
singleWarning (forall a. Located a => a -> SrcLoc
srclocOf loc
loc) Doc ()
problem
        }

  newName :: VName -> TypeM VName
newName VName
v = do
    TypeState
s <- forall s (m :: * -> *). MonadState s m => m s
get
    let (VName
v', VNameSource
src') = VNameSource -> VName -> (VName, VNameSource)
Futhark.FreshNames.newName (TypeState -> VNameSource
stateNameSource TypeState
s) VName
v
    forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ TypeState
s {stateNameSource :: VNameSource
stateNameSource = VNameSource
src'}
    forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v'

  newID :: Name -> TypeM VName
newID Name
s = forall (m :: * -> *). MonadTypeChecker m => VName -> m VName
newName forall a b. (a -> b) -> a -> b
$ Name -> Int -> VName
VName Name
s Int
0

  newTypeName :: Name -> TypeM VName
newTypeName Name
name = do
    Int
i <- TypeM Int
incCounter
    forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID forall a b. (a -> b) -> a -> b
$ Name -> Int -> Name
mkTypeVarName Name
name Int
i

  bindNameMap :: forall a.
Map (Namespace, Name) (QualName VName) -> TypeM a -> TypeM a
bindNameMap Map (Namespace, Name) (QualName VName)
m = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
    let env :: Env
env = Context -> Env
contextEnv Context
ctx
     in Context
ctx {contextEnv :: Env
contextEnv = Env
env {envNameMap :: Map (Namespace, Name) (QualName VName)
envNameMap = Map (Namespace, Name) (QualName VName)
m forall a. Semigroup a => a -> a -> a
<> Env -> Map (Namespace, Name) (QualName VName)
envNameMap Env
env}}

  bindVal :: forall a. VName -> BoundV -> TypeM a -> TypeM a
bindVal VName
v BoundV
t = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
    Context
ctx
      { contextEnv :: Env
contextEnv =
          (Context -> Env
contextEnv Context
ctx)
            { envVtable :: Map VName BoundV
envVtable = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v BoundV
t forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable forall a b. (a -> b) -> a -> b
$ Context -> Env
contextEnv Context
ctx
            }
      }

  checkQualName :: Namespace -> QualName Name -> SrcLoc -> TypeM (QualName VName)
checkQualName Namespace
space QualName Name
name SrcLoc
loc = forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv Namespace
space QualName Name
name SrcLoc
loc

  lookupType :: SrcLoc
-> QualName Name
-> TypeM (QualName VName, [TypeParam], StructRetType, Liftedness)
lookupType SrcLoc
loc QualName Name
qn = do
    Env
outer_env <- TypeM Env
askEnv
    (Env
scope, qn' :: QualName VName
qn'@(QualName [VName]
qs VName
name)) <- Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv Namespace
Type QualName Name
qn SrcLoc
loc
    case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall a b. (a -> b) -> a -> b
$ Env -> Map VName TypeBinding
envTypeTable Env
scope of
      Maybe TypeBinding
Nothing -> forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
unknownType SrcLoc
loc QualName Name
qn
      Just (TypeAbbr Liftedness
l [TypeParam]
ps (RetType [VName]
dims TypeBase Size ()
def)) ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName
qn', [TypeParam]
ps, forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims forall a b. (a -> b) -> a -> b
$ forall as.
Env -> [VName] -> [VName] -> TypeBase Size as -> TypeBase Size as
qualifyTypeVars Env
outer_env forall a. Monoid a => a
mempty [VName]
qs TypeBase Size ()
def, Liftedness
l)

  lookupMod :: SrcLoc -> QualName Name -> TypeM (QualName VName, Mod)
lookupMod SrcLoc
loc QualName Name
qn = do
    (Env
scope, qn' :: QualName VName
qn'@(QualName [VName]
_ VName
name)) <- Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv Namespace
Term QualName Name
qn SrcLoc
loc
    case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall a b. (a -> b) -> a -> b
$ Env -> Map VName Mod
envModTable Env
scope of
      Maybe Mod
Nothing -> forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
Term QualName Name
qn SrcLoc
loc
      Just Mod
m -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName
qn', Mod
m)

  lookupVar :: SrcLoc -> QualName Name -> TypeM (QualName VName, PatType)
lookupVar SrcLoc
loc QualName Name
qn = do
    Env
outer_env <- TypeM Env
askEnv
    (Env
env, qn' :: QualName VName
qn'@(QualName [VName]
qs VName
name)) <- Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv Namespace
Term QualName Name
qn SrcLoc
loc
    case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable Env
env of
      Maybe BoundV
Nothing -> forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
Term QualName Name
qn SrcLoc
loc
      Just (BoundV [TypeParam]
_ TypeBase Size ()
t)
        | String
"_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` VName -> String
baseString VName
name -> forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
underscoreUse SrcLoc
loc QualName Name
qn
        | Bool
otherwise ->
            case forall dim as. TypeBase dim as -> Maybe (TypeBase dim as)
getType TypeBase Size ()
t of
              Maybe (TypeBase Size ())
Nothing ->
                forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                  Doc ()
"Attempt to use function" forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall v a. IsName v => v -> Doc a
prettyName VName
name forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
"as value."
              Just TypeBase Size ()
t' ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure
                  ( QualName VName
qn',
                    forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct forall a b. (a -> b) -> a -> b
$
                      forall as.
Env -> [VName] -> [VName] -> TypeBase Size as -> TypeBase Size as
qualifyTypeVars Env
outer_env forall a. Monoid a => a
mempty [VName]
qs TypeBase Size ()
t'
                  )

  typeError :: forall loc a. Located loc => loc -> Notes -> Doc () -> TypeM a
typeError loc
loc Notes
notes Doc ()
s = forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$ Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes Doc ()
s

-- | Extract from a type a first-order type.
getType :: TypeBase dim as -> Maybe (TypeBase dim as)
getType :: forall dim as. TypeBase dim as -> Maybe (TypeBase dim as)
getType (Scalar Arrow {}) = forall a. Maybe a
Nothing
getType TypeBase dim as
t = forall a. a -> Maybe a
Just TypeBase dim as
t

checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName)
checkQualNameWithEnv Namespace
space qn :: QualName Name
qn@(QualName [Name]
quals Name
name) SrcLoc
loc = do
  Env
env <- TypeM Env
askEnv
  forall {f :: * -> *}.
MonadTypeChecker f =>
Env -> [Name] -> f (Env, QualName VName)
descend Env
env [Name]
quals
  where
    descend :: Env -> [Name] -> f (Env, QualName VName)
descend Env
scope []
      | Just QualName VName
name' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
space, Name
name) forall a b. (a -> b) -> a -> b
$ Env -> Map (Namespace, Name) (QualName VName)
envNameMap Env
scope =
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env
scope, QualName VName
name')
      | Bool
otherwise =
          forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
qn SrcLoc
loc
    descend Env
scope (Name
q : [Name]
qs)
      | Just (QualName [VName]
_ VName
q') <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
Term, Name
q) forall a b. (a -> b) -> a -> b
$ Env -> Map (Namespace, Name) (QualName VName)
envNameMap Env
scope,
        Just Mod
res <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q' forall a b. (a -> b) -> a -> b
$ Env -> Map VName Mod
envModTable Env
scope =
          case Mod
res of
            ModEnv Env
q_scope -> do
              (Env
scope', QualName [VName]
qs' VName
name') <- Env -> [Name] -> f (Env, QualName VName)
descend Env
q_scope [Name]
qs
              forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env
scope', forall vn. [vn] -> vn -> QualName vn
QualName (VName
q' forall a. a -> [a] -> [a]
: [VName]
qs') VName
name')
            ModFun {} -> forall (m :: * -> *) a. MonadTypeChecker m => SrcLoc -> m a
unappliedFunctor SrcLoc
loc
      | Bool
otherwise =
          forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
qn SrcLoc
loc

-- | 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 :: forall as.
Env -> [VName] -> [VName] -> TypeBase Size as -> TypeBase Size as
qualifyTypeVars Env
outer_env [VName]
orig_except [VName]
ref_qs = forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType (forall a. Ord a => [a] -> Set a
S.fromList [VName]
orig_except)
  where
    onType ::
      S.Set VName ->
      TypeBase Size as ->
      TypeBase Size as
    onType :: forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except (Array as
as Uniqueness
u Shape Size
shape ScalarTypeBase Size ()
et) =
      forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array as
as Uniqueness
u (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall {t :: * -> *}. Foldable t => t VName -> Size -> Size
onDim Set VName
except) Shape Size
shape) (forall {as}.
Set VName -> ScalarTypeBase Size as -> ScalarTypeBase Size as
onScalar Set VName
except ScalarTypeBase Size ()
et)
    onType Set VName
except (Scalar ScalarTypeBase Size as
t) =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall {as}.
Set VName -> ScalarTypeBase Size as -> ScalarTypeBase Size as
onScalar Set VName
except ScalarTypeBase Size as
t

    onScalar :: Set VName -> ScalarTypeBase Size as -> ScalarTypeBase Size as
onScalar Set VName
_ (Prim PrimType
t) = forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    onScalar Set VName
except (TypeVar as
as Uniqueness
u QualName VName
qn [TypeArg Size]
targs) =
      forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar as
as Uniqueness
u (forall {t :: * -> *}.
Foldable t =>
t VName -> QualName VName -> QualName VName
qual Set VName
except QualName VName
qn) (forall a b. (a -> b) -> [a] -> [b]
map (Set VName -> TypeArg Size -> TypeArg Size
onTypeArg Set VName
except) [TypeArg Size]
targs)
    onScalar Set VName
except (Record Map Name (TypeBase Size as)
m) =
      forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except) Map Name (TypeBase Size as)
m
    onScalar Set VName
except (Sum Map Name [TypeBase Size as]
m) =
      forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a -> b) -> a -> b
$ forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except) Map Name [TypeBase Size as]
m
    onScalar Set VName
except (Arrow as
as PName
p TypeBase Size ()
t1 (RetType [VName]
dims TypeBase Size as
t2)) =
      forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
as PName
p (forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except' TypeBase Size ()
t1) forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except' TypeBase Size as
t2)
      where
        except' :: Set VName
except' = case PName
p of
          Named VName
p' -> forall a. Ord a => a -> Set a -> Set a
S.insert VName
p' Set VName
except
          PName
Unnamed -> Set VName
except

    onTypeArg :: Set VName -> TypeArg Size -> TypeArg Size
onTypeArg Set VName
except (TypeArgDim Size
d SrcLoc
loc) =
      forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim (forall {t :: * -> *}. Foldable t => t VName -> Size -> Size
onDim Set VName
except Size
d) SrcLoc
loc
    onTypeArg Set VName
except (TypeArgType TypeBase Size ()
t SrcLoc
loc) =
      forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except TypeBase Size ()
t) SrcLoc
loc

    onDim :: t VName -> Size -> Size
onDim t VName
except (NamedSize QualName VName
qn) = QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall {t :: * -> *}.
Foldable t =>
t VName -> QualName VName -> QualName VName
qual t VName
except QualName VName
qn
    onDim t VName
_ Size
d = Size
d

    qual :: t VName -> QualName VName -> QualName VName
qual t VName
except (QualName [VName]
orig_qs VName
name)
      | VName
name forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
except Bool -> Bool -> Bool
|| [VName] -> VName -> Env -> Bool
reachable [VName]
orig_qs VName
name Env
outer_env =
          forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name
      | Bool
otherwise =
          [VName] -> [VName] -> QualName VName -> QualName VName
prependAsNecessary [] [VName]
ref_qs forall a b. (a -> b) -> a -> b
$ forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name

    prependAsNecessary :: [VName] -> [VName] -> QualName VName -> QualName VName
prependAsNecessary [VName]
qs [VName]
rem_qs (QualName [VName]
orig_qs VName
name)
      | [VName] -> VName -> Env -> Bool
reachable ([VName]
qs forall a. [a] -> [a] -> [a]
++ [VName]
orig_qs) VName
name Env
outer_env = forall vn. [vn] -> vn -> QualName vn
QualName ([VName]
qs forall a. [a] -> [a] -> [a]
++ [VName]
orig_qs) VName
name
      | Bool
otherwise = case [VName]
rem_qs of
          VName
q : [VName]
rem_qs' -> [VName] -> [VName] -> QualName VName -> QualName VName
prependAsNecessary ([VName]
qs forall a. [a] -> [a] -> [a]
++ [VName
q]) [VName]
rem_qs' (forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name)
          [] -> forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name

    reachable :: [VName] -> VName -> Env -> Bool
reachable [] VName
name Env
env =
      VName
name forall k a. Ord k => k -> Map k a -> Bool
`M.member` Env -> Map VName BoundV
envVtable Env
env
        Bool -> Bool -> Bool
|| forall a. Maybe a -> Bool
isJust (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find TypeBinding -> Bool
matches forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
M.elems (Env -> Map VName TypeBinding
envTypeTable Env
env))
      where
        matches :: TypeBinding -> Bool
matches (TypeAbbr Liftedness
_ [TypeParam]
_ (RetType [VName]
_ (Scalar (TypeVar ()
_ Uniqueness
_ (QualName [VName]
x_qs VName
name') [TypeArg Size]
_)))) =
          forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
x_qs Bool -> Bool -> Bool
&& VName
name forall a. Eq a => a -> a -> Bool
== VName
name'
        matches TypeBinding
_ = Bool
False
    reachable (VName
q : [VName]
qs') VName
name Env
env
      | Just (ModEnv Env
env') <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q forall a b. (a -> b) -> a -> b
$ Env -> Map VName Mod
envModTable Env
env =
          [VName] -> VName -> Env -> Bool
reachable [VName]
qs' VName
name Env
env'
      | Bool
otherwise = Bool
False

-- | Turn a 'Left' 'TypeError' into an actual error.
badOnLeft :: Either TypeError a -> TypeM a
badOnLeft :: forall a. Either TypeError a -> TypeM a
badOnLeft = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | All signed integer types.
anySignedType :: [PrimType]
anySignedType :: [PrimType]
anySignedType = forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]

-- | All unsigned integer types.
anyUnsignedType :: [PrimType]
anyUnsignedType :: [PrimType]
anyUnsignedType = forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]

-- | All integer types.
anyIntType :: [PrimType]
anyIntType :: [PrimType]
anyIntType = [PrimType]
anySignedType forall a. [a] -> [a] -> [a]
++ [PrimType]
anyUnsignedType

-- | All floating-point types.
anyFloatType :: [PrimType]
anyFloatType :: [PrimType]
anyFloatType = forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]

-- | All number types.
anyNumberType :: [PrimType]
anyNumberType :: [PrimType]
anyNumberType = [PrimType]
anyIntType forall a. [a] -> [a] -> [a]
++ [PrimType]
anyFloatType

-- | All primitive types.
anyPrimType :: [PrimType]
anyPrimType :: [PrimType]
anyPrimType = PrimType
Bool forall a. a -> [a] -> [a]
: [PrimType]
anyIntType forall a. [a] -> [a] -> [a]
++ [PrimType]
anyFloatType

--- Name handling

-- | The 'NameMap' corresponding to the intrinsics module.
intrinsicsNameMap :: NameMap
intrinsicsNameMap :: Map (Namespace, Name) (QualName VName)
intrinsicsNameMap = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (VName, Intrinsic) -> ((Namespace, Name), QualName VName)
mapping forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map VName Intrinsic
intrinsics
  where
    mapping :: (VName, Intrinsic) -> ((Namespace, Name), QualName VName)
mapping (VName
v, IntrinsicType {}) = ((Namespace
Type, VName -> Name
baseName VName
v), forall vn. [vn] -> vn -> QualName vn
QualName [] VName
v)
    mapping (VName
v, Intrinsic
_) = ((Namespace
Term, VName -> Name
baseName VName
v), forall vn. [vn] -> vn -> QualName vn
QualName [] VName
v)

-- | The names that are available in the initial environment.
topLevelNameMap :: NameMap
topLevelNameMap :: Map (Namespace, Name) (QualName VName)
topLevelNameMap = forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\(Namespace, Name)
k QualName VName
_ -> (Namespace, Name) -> Bool
available (Namespace, Name)
k) Map (Namespace, Name) (QualName VName)
intrinsicsNameMap
  where
    available :: (Namespace, Name) -> Bool
    available :: (Namespace, Name) -> Bool
available (Namespace
Type, Name
_) = Bool
True
    available (Namespace
Term, Name
v) = Name
v forall a. Ord a => a -> Set a -> Bool
`S.member` (Set Name
type_names forall a. Semigroup a => a -> a -> a
<> Set Name
binop_names forall a. Semigroup a => a -> a -> a
<> Set Name
fun_names)
      where
        type_names :: Set Name
type_names = forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Text -> Name
nameFromText forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> Text
prettyText) [PrimType]
anyPrimType
        binop_names :: Set Name
binop_names =
          forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$
            forall a b. (a -> b) -> [a] -> [b]
map
              (Text -> Name
nameFromText forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> Text
prettyText)
              [forall a. Bounded a => a
minBound .. (forall a. Bounded a => a
maxBound :: BinOp)]
        fun_names :: Set Name
fun_names = forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map String -> Name
nameFromString [String
"shape"]
    available (Namespace, Name)
_ = Bool
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 :: Name -> Int -> Name
mkTypeVarName Name
desc Int
i =
  Name
desc forall a. Semigroup a => a -> a -> a
<> String -> Name
nameFromString (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Char -> Maybe Char
subscript (forall a. Show a => a -> String
show Int
i))
  where
    subscript :: Char -> Maybe Char
subscript = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip String
"0123456789" String
"₀₁₂₃₄₅₆₇₈₉"

-- | Type-check an attribute.
checkAttr :: MonadTypeChecker m => AttrInfo Name -> m (AttrInfo VName)
checkAttr :: forall (m :: * -> *).
MonadTypeChecker m =>
AttrInfo Name -> m (AttrInfo VName)
checkAttr (AttrComp Name
f [AttrInfo Name]
attrs SrcLoc
loc) =
  forall {k} (vn :: k).
Name -> [AttrInfo vn] -> SrcLoc -> AttrInfo vn
AttrComp Name
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *).
MonadTypeChecker m =>
AttrInfo Name -> m (AttrInfo VName)
checkAttr [AttrInfo Name]
attrs forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkAttr (AttrAtom (AtomName Name
v) SrcLoc
loc) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (vn :: k). AttrAtom vn -> SrcLoc -> AttrInfo vn
AttrAtom (forall {k} (vn :: k). Name -> AttrAtom vn
AtomName Name
v) SrcLoc
loc
checkAttr (AttrAtom (AtomInt Integer
x) SrcLoc
loc) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (vn :: k). AttrAtom vn -> SrcLoc -> AttrInfo vn
AttrAtom (forall {k} (vn :: k). Integer -> AttrAtom vn
AtomInt Integer
x) SrcLoc
loc