-- | 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 (NonEmpty Notes -> Notes
Notes -> Notes -> Notes
(Notes -> Notes -> Notes)
-> (NonEmpty Notes -> Notes)
-> (forall b. Integral b => b -> Notes -> Notes)
-> Semigroup 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
$c<> :: Notes -> Notes -> Notes
<> :: Notes -> Notes -> Notes
$csconcat :: NonEmpty Notes -> Notes
sconcat :: NonEmpty Notes -> Notes
$cstimes :: forall b. Integral b => b -> Notes -> Notes
stimes :: forall b. Integral b => b -> Notes -> Notes
Semigroup, Semigroup Notes
Notes
Semigroup Notes =>
Notes
-> (Notes -> Notes -> Notes) -> ([Notes] -> Notes) -> Monoid Notes
[Notes] -> Notes
Notes -> Notes -> Notes
forall a.
Semigroup a =>
a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
$cmempty :: Notes
mempty :: Notes
$cmappend :: Notes -> Notes -> Notes
mappend :: Notes -> Notes -> Notes
$cmconcat :: [Notes] -> Notes
mconcat :: [Notes] -> Notes
Monoid)

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

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

-- | A single note.
aNote :: Doc () -> Notes
aNote :: Doc () -> Notes
aNote = [Note] -> Notes
Notes ([Note] -> Notes) -> (Doc () -> [Note]) -> Doc () -> Notes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Note -> [Note]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Note -> [Note]) -> (Doc () -> Note) -> Doc () -> [Note]
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) =
  AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall ann. ann -> Doc ann -> Doc ann
annotate
    (AnsiStyle
bold AnsiStyle -> AnsiStyle -> AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Color -> AnsiStyle
color Color
Red)
    (Doc AnsiStyle
"Error at " Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Text -> Doc AnsiStyle
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (SrcLoc -> Text
forall a. Located a => a -> Text
locText (Loc -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Loc
loc)) Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Doc AnsiStyle
":")
    Doc AnsiStyle -> Doc AnsiStyle -> 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) =
  Doc () -> Doc AnsiStyle
forall ann xxx. Doc ann -> Doc xxx
unAnnotate Doc ()
msg Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Notes -> Doc AnsiStyle
forall a ann. Pretty a => a -> Doc ann
forall ann. Notes -> Doc ann
pretty Notes
notes Doc AnsiStyle -> Doc AnsiStyle -> Doc AnsiStyle
forall a. Semigroup a => a -> a -> a
<> Doc AnsiStyle
forall ann. Doc ann
hardline

errorIndexUrl :: Doc a
errorIndexUrl :: forall ann. Doc ann
errorIndexUrl = Doc a
version_url Doc a -> Doc a -> Doc a
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
      | [Int] -> Int
forall a. HasCallStack => [a] -> a
last (Version -> [Int]
Version.versionBranch Version
version) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Doc a
base_url Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"latest/"
      | Bool
otherwise = Doc a
base_url Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"v" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> String -> Doc a
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Version -> String
Version.showVersion Version
version) Doc a -> Doc a -> Doc a
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 =
  [Doc a] -> Doc a
forall a. [Doc a] -> Doc a
stack
    [ Doc a
msg,
      Doc a
"\nFor more information, see:",
      Int -> Doc a -> Doc a
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc a
forall ann. Doc ann
errorIndexUrl Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"#" Doc a -> Doc 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 =
  SrcLoc -> Notes -> Doc () -> m a
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty Doc ()
"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 =
  SrcLoc -> Notes -> Doc () -> m a
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> m a) -> Doc () -> m a
forall a b. (a -> b) -> a -> b
$
    Doc ()
"Unknown" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Namespace -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Namespace -> Doc ann
pretty Namespace
space Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (QualName Name -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName Name -> Doc ann
pretty QualName Name
name)

-- | 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 =
  SrcLoc -> Notes -> Doc () -> m a
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> m a) -> Doc () -> m a
forall a b. (a -> b) -> a -> b
$
    Doc ()
"Use of"
      Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (QualName Name -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName Name -> Doc ann
pretty QualName Name
name)
      Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
": 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
  { 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,
    -- | Which names have been used.
    TypeState -> Set VName
stateUsed :: S.Set VName,
    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
Applicative TypeM =>
(forall a b. TypeM a -> (a -> TypeM b) -> TypeM b)
-> (forall a b. TypeM a -> TypeM b -> TypeM b)
-> (forall a. a -> TypeM a)
-> Monad 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
$c>>= :: forall a b. TypeM a -> (a -> TypeM b) -> TypeM b
>>= :: forall a b. TypeM a -> (a -> TypeM b) -> TypeM b
$c>> :: forall a b. TypeM a -> TypeM b -> TypeM b
>> :: forall a b. TypeM a -> TypeM b -> TypeM b
$creturn :: forall a. a -> TypeM a
return :: forall a. a -> TypeM a
Monad,
      (forall a b. (a -> b) -> TypeM a -> TypeM b)
-> (forall a b. a -> TypeM b -> TypeM a) -> Functor TypeM
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
$cfmap :: forall a b. (a -> b) -> TypeM a -> TypeM b
fmap :: forall a b. (a -> b) -> TypeM a -> TypeM b
$c<$ :: forall a b. a -> TypeM b -> TypeM a
<$ :: forall a b. a -> TypeM b -> TypeM a
Functor,
      Functor TypeM
Functor TypeM =>
(forall a. a -> TypeM a)
-> (forall a b. TypeM (a -> b) -> TypeM a -> TypeM b)
-> (forall a b c. (a -> b -> c) -> TypeM a -> TypeM b -> TypeM c)
-> (forall a b. TypeM a -> TypeM b -> TypeM b)
-> (forall a b. TypeM a -> TypeM b -> TypeM a)
-> Applicative 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
$cpure :: forall a. a -> TypeM a
pure :: forall a. a -> TypeM a
$c<*> :: forall a b. TypeM (a -> b) -> TypeM a -> TypeM b
<*> :: forall a b. TypeM (a -> b) -> TypeM a -> TypeM b
$cliftA2 :: forall a b c. (a -> b -> c) -> TypeM a -> TypeM b -> TypeM c
liftA2 :: forall a b c. (a -> b -> c) -> TypeM a -> TypeM b -> TypeM c
$c*> :: forall a b. TypeM a -> TypeM b -> TypeM b
*> :: forall a b. TypeM a -> TypeM b -> TypeM b
$c<* :: forall a b. TypeM a -> TypeM b -> TypeM a
<* :: forall a b. TypeM a -> TypeM b -> TypeM a
Applicative,
      MonadReader Context,
      MonadState TypeState
    )

instance MonadError TypeError TypeM where
  throwError :: forall a. TypeError -> TypeM a
throwError TypeError
e = ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
forall a.
ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
TypeM (ReaderT
   Context (StateT TypeState (Except (Warnings, TypeError))) a
 -> TypeM a)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
forall a b. (a -> b) -> a -> b
$ do
    Warnings
ws <- (TypeState -> Warnings)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) Warnings
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TypeState -> Warnings
stateWarnings
    (Warnings, TypeError)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
forall a.
(Warnings, TypeError)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
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 =
    ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
forall a.
ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
TypeM (ReaderT
   Context (StateT TypeState (Except (Warnings, TypeError))) a
 -> TypeM a)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
-> TypeM a
forall a b. (a -> b) -> a -> b
$ ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
m ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> ((Warnings, TypeError)
    -> ReaderT
         Context (StateT TypeState (Except (Warnings, TypeError))) a)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
forall a.
ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> ((Warnings, TypeError)
    -> ReaderT
         Context (StateT TypeState (Except (Warnings, TypeError))) a)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
forall e (m :: * -> *) a.
MonadError e m =>
m a -> (e -> m a) -> m a
`catchError` (Warnings, TypeError)
-> ReaderT
     Context (StateT TypeState (Except (Warnings, TypeError))) a
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 -> Set VName -> Int -> TypeState
TypeState VNameSource
src Warnings
forall a. Monoid a => a
mempty Set VName
forall a. Monoid a => a
mempty Int
0
  case Except (Warnings, TypeError) (a, TypeState)
-> Either (Warnings, TypeError) (a, TypeState)
forall e a. Except e a -> Either e a
runExcept (Except (Warnings, TypeError) (a, TypeState)
 -> Either (Warnings, TypeError) (a, TypeState))
-> Except (Warnings, TypeError) (a, TypeState)
-> Either (Warnings, TypeError) (a, TypeState)
forall a b. (a -> b) -> a -> b
$ StateT TypeState (Except (Warnings, TypeError)) a
-> TypeState -> Except (Warnings, TypeError) (a, TypeState)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ReaderT Context (StateT TypeState (Except (Warnings, TypeError))) a
-> Context -> StateT TypeState (Except (Warnings, TypeError)) a
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, TypeError -> Either TypeError (a, VNameSource)
forall a b. a -> Either a b
Left TypeError
e)
    Right (a
x, TypeState
s') -> (TypeState -> Warnings
stateWarnings TypeState
s', (a, VNameSource) -> Either TypeError (a, VNameSource)
forall a b. b -> Either a b
Right (a
x, TypeState -> VNameSource
stateNameSource TypeState
s'))

-- | Retrieve the current 'Env'.
askEnv :: TypeM Env
askEnv :: TypeM Env
askEnv = (Context -> Env) -> TypeM Env
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 = (Context -> ImportName) -> TypeM ImportName
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 = (Context -> Bool) -> TypeM Bool
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 = (Context -> Context) -> TypeM a -> TypeM a
forall a. (Context -> Context) -> TypeM a -> TypeM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Context -> Context) -> TypeM a -> TypeM a)
-> (Context -> Context) -> TypeM a -> TypeM a
forall a b. (a -> b) -> a -> b
$ \Context
ctx -> Context
ctx {contextAtTopLevel = 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',) (MTy -> (QualName VName, MTy))
-> TypeM MTy -> TypeM (QualName VName, MTy)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeM MTy -> (MTy -> TypeM MTy) -> Maybe MTy -> TypeM MTy
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeM MTy
forall {a}. TypeM a
explode MTy -> TypeM MTy
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Map VName MTy -> Maybe MTy
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName MTy -> Maybe MTy) -> Map VName MTy -> Maybe MTy
forall a b. (a -> b) -> a -> b
$ Env -> Map VName MTy
envModTypeTable Env
scope)
  where
    explode :: TypeM a
explode = Namespace -> QualName Name -> SrcLoc -> TypeM a
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 (ImportName, Env)
lookupImport :: SrcLoc -> String -> TypeM (ImportName, Env)
lookupImport SrcLoc
loc String
file = do
  ImportTable
imports <- (Context -> ImportTable) -> TypeM ImportTable
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> ImportTable
contextImportTable
  ImportName
my_path <- (Context -> ImportName) -> TypeM ImportName
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Context -> ImportName
contextImportName
  let canonical_import :: ImportName
canonical_import = ImportName -> String -> ImportName
mkImportFrom ImportName
my_path String
file
  case ImportName -> ImportTable -> Maybe Env
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ImportName
canonical_import ImportTable
imports of
    Maybe Env
Nothing ->
      SrcLoc -> Notes -> Doc () -> TypeM (ImportName, Env)
forall loc a. Located loc => loc -> Notes -> Doc () -> TypeM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> TypeM (ImportName, Env))
-> Doc () -> TypeM (ImportName, Env)
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Unknown import"
          Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Text -> Doc ()
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (ImportName -> Text
includeToText ImportName
canonical_import))
          Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"Known:"
          Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
commasep ((ImportName -> Doc ()) -> [ImportName] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Doc ()
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Text -> Doc ()) -> (ImportName -> Text) -> ImportName -> Doc ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ImportName -> Text
includeToText) (ImportTable -> [ImportName]
forall k a. Map k a -> [k]
M.keys ImportTable
imports))
    Just Env
scope -> (ImportName, Env) -> TypeM (ImportName, Env)
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ImportName
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 = (Context -> Context) -> TypeM a -> TypeM a
forall a. (Context -> Context) -> TypeM a -> TypeM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Context -> Context) -> TypeM a -> TypeM a)
-> (Context -> Context) -> TypeM a -> TypeM a
forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
  let env' :: Env
env' = Env
env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Context -> Env
contextEnv Context
ctx
   in Context
ctx {contextEnv = env'}

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

bindNameMap :: NameMap -> TypeM a -> TypeM a
bindNameMap :: forall a. NameMap -> TypeM a -> TypeM a
bindNameMap NameMap
m = (Context -> Context) -> TypeM a -> TypeM a
forall a. (Context -> Context) -> TypeM a -> TypeM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Context -> Context) -> TypeM a -> TypeM a)
-> (Context -> Context) -> TypeM a -> TypeM a
forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
  let env :: Env
env = Context -> Env
contextEnv Context
ctx
   in Context
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 Name
s = VName -> m VName
forall (m :: * -> *). MonadTypeChecker m => VName -> m VName
newName (VName -> m VName) -> VName -> m VName
forall a b. (a -> b) -> a -> b
$ Name -> Int -> VName
VName Name
s Int
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 :: (Namespace, VName, SrcLoc) -> TypeM ()
warnIfUnused (Namespace
ns, VName
name, SrcLoc
loc) = do
  Set VName
used <- (TypeState -> Set VName) -> TypeM (Set VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TypeState -> Set VName
stateUsed
  Bool -> TypeM () -> TypeM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName
name VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
used Bool -> Bool -> Bool
|| Text
"_" Text -> Text -> Bool
`T.isPrefixOf` Name -> Text
nameToText (VName -> Name
baseName VName
name)) (TypeM () -> TypeM ()) -> TypeM () -> TypeM ()
forall a b. (a -> b) -> a -> b
$
    SrcLoc -> Doc () -> TypeM ()
forall loc. Located loc => loc -> Doc () -> TypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn SrcLoc
loc (Doc () -> TypeM ()) -> Doc () -> TypeM ()
forall a b. (a -> b) -> a -> b
$
      Doc ()
"Unused" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Namespace -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Namespace -> Doc ann
pretty Namespace
ns Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
name) Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."

-- | 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 :: forall a.
[(Namespace, Name, SrcLoc)] -> ([VName] -> TypeM a) -> TypeM a
bindSpaced [(Namespace, Name, SrcLoc)]
names [VName] -> TypeM a
body = do
  [VName]
names' <- ((Namespace, Name, SrcLoc) -> TypeM VName)
-> [(Namespace, Name, SrcLoc)] -> TypeM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\(Namespace
_, Name
v, SrcLoc
_) -> Name -> TypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID Name
v) [(Namespace, Name, SrcLoc)]
names
  let mapping :: NameMap
mapping = [((Namespace, Name), QualName VName)] -> NameMap
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([((Namespace, Name), QualName VName)] -> NameMap)
-> [((Namespace, Name), QualName VName)] -> NameMap
forall a b. (a -> b) -> a -> b
$ [(Namespace, Name)]
-> [QualName VName] -> [((Namespace, Name), QualName VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Namespace, Name, SrcLoc) -> (Namespace, Name))
-> [(Namespace, Name, SrcLoc)] -> [(Namespace, Name)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Namespace
ns, Name
v, SrcLoc
_) -> (Namespace
ns, Name
v)) [(Namespace, Name, SrcLoc)]
names) ([QualName VName] -> [((Namespace, Name), QualName VName)])
-> [QualName VName] -> [((Namespace, Name), QualName VName)]
forall a b. (a -> b) -> a -> b
$ (VName -> QualName VName) -> [VName] -> [QualName VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> QualName VName
forall v. v -> QualName v
qualName [VName]
names'
  NameMap -> TypeM a -> TypeM a
forall a. NameMap -> TypeM a -> TypeM a
bindNameMap NameMap
mapping ([VName] -> TypeM a
body [VName]
names')
    TypeM a -> TypeM () -> TypeM a
forall a b. TypeM a -> TypeM b -> TypeM a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ((Namespace, VName, SrcLoc) -> TypeM ())
-> [(Namespace, VName, SrcLoc)] -> TypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Namespace, VName, SrcLoc) -> TypeM ()
warnIfUnused [(Namespace
ns, VName
v, SrcLoc
loc) | ((Namespace
ns, Name
_, SrcLoc
loc), VName
v) <- [(Namespace, Name, SrcLoc)]
-> [VName] -> [((Namespace, Name, SrcLoc), VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Namespace, Name, SrcLoc)]
names [VName]
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 :: forall a.
Namespace -> Name -> SrcLoc -> (VName -> TypeM a) -> TypeM a
bindSpaced1 Namespace
ns Name
name SrcLoc
loc VName -> TypeM a
body = do
  VName
name' <- Name -> TypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID Name
name
  let mapping :: NameMap
mapping = (Namespace, Name) -> QualName VName -> NameMap
forall k a. k -> a -> Map k a
M.singleton (Namespace
ns, Name
name) (QualName VName -> NameMap) -> QualName VName -> NameMap
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
name'
  NameMap -> TypeM a -> TypeM a
forall a. NameMap -> TypeM a -> TypeM a
bindNameMap NameMap
mapping (VName -> TypeM a
body VName
name') TypeM a -> TypeM () -> TypeM a
forall a b. TypeM a -> TypeM b -> TypeM a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* (Namespace, VName, SrcLoc) -> TypeM ()
warnIfUnused (Namespace
ns, VName
name', SrcLoc
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 :: forall {k} (t :: k) a.
[IdentBase NoInfo VName t] -> TypeM a -> TypeM a
bindIdents [IdentBase NoInfo VName t]
idents TypeM a
body = do
  let mapping :: NameMap
mapping =
        [((Namespace, Name), QualName VName)] -> NameMap
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([((Namespace, Name), QualName VName)] -> NameMap)
-> [((Namespace, Name), QualName VName)] -> NameMap
forall a b. (a -> b) -> a -> b
$
          [(Namespace, Name)]
-> [QualName VName] -> [((Namespace, Name), QualName VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip
            ((IdentBase NoInfo VName t -> (Namespace, Name))
-> [IdentBase NoInfo VName t] -> [(Namespace, Name)]
forall a b. (a -> b) -> [a] -> [b]
map ((Namespace
Term,) (Name -> (Namespace, Name))
-> (IdentBase NoInfo VName t -> Name)
-> IdentBase NoInfo VName t
-> (Namespace, Name)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Name
baseName (VName -> Name)
-> (IdentBase NoInfo VName t -> VName)
-> IdentBase NoInfo VName t
-> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentBase NoInfo VName t -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName)) [IdentBase NoInfo VName t]
idents)
            ((IdentBase NoInfo VName t -> QualName VName)
-> [IdentBase NoInfo VName t] -> [QualName VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> QualName VName)
-> (IdentBase NoInfo VName t -> VName)
-> IdentBase NoInfo VName t
-> QualName VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentBase NoInfo VName t -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName) [IdentBase NoInfo VName t]
idents)
  NameMap -> TypeM a -> TypeM a
forall a. NameMap -> TypeM a -> TypeM a
bindNameMap NameMap
mapping TypeM a
body TypeM a -> TypeM () -> TypeM a
forall a b. TypeM a -> TypeM b -> TypeM a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ((Namespace, VName, SrcLoc) -> TypeM ())
-> [(Namespace, VName, SrcLoc)] -> TypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Namespace, VName, SrcLoc) -> TypeM ()
warnIfUnused [(Namespace
Term, VName
v, SrcLoc
loc) | Ident VName
v NoInfo t
_ SrcLoc
loc <- [IdentBase NoInfo VName t]
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 :: VName -> TypeM ()
usedName VName
name = (TypeState -> TypeState) -> TypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TypeState -> TypeState) -> TypeM ())
-> (TypeState -> TypeState) -> TypeM ()
forall a b. (a -> b) -> a -> b
$ \TypeState
s -> TypeState
s {stateUsed = S.insert name $ stateUsed s}

instance MonadTypeChecker TypeM where
  warnings :: Warnings -> TypeM ()
warnings Warnings
ws =
    (TypeState -> TypeState) -> TypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TypeState -> TypeState) -> TypeM ())
-> (TypeState -> TypeState) -> TypeM ()
forall a b. (a -> b) -> a -> b
$ \TypeState
s -> TypeState
s {stateWarnings = stateWarnings s <> ws}

  warn :: forall loc. Located loc => loc -> Doc () -> TypeM ()
warn loc
loc Doc ()
problem =
    Warnings -> TypeM ()
forall (m :: * -> *). MonadTypeChecker m => Warnings -> m ()
warnings (Warnings -> TypeM ()) -> Warnings -> TypeM ()
forall a b. (a -> b) -> a -> b
$ Loc -> Doc () -> Warnings
singleWarning (loc -> Loc
forall a. Located a => a -> Loc
locOf loc
loc) Doc ()
problem

  newName :: VName -> TypeM VName
newName VName
v = do
    TypeState
s <- TypeM TypeState
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
    TypeState -> TypeM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (TypeState -> TypeM ()) -> TypeState -> TypeM ()
forall a b. (a -> b) -> a -> b
$ TypeState
s {stateNameSource = src'}
    VName -> TypeM VName
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v'

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

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

  lookupType :: QualName VName -> TypeM ([TypeParam], StructRetType, Liftedness)
lookupType QualName VName
qn = do
    Env
outer_env <- TypeM Env
askEnv
    Env
scope <- QualName VName -> TypeM Env
lookupQualNameEnv QualName VName
qn
    case VName -> Map VName TypeBinding -> Maybe TypeBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn) (Map VName TypeBinding -> Maybe TypeBinding)
-> Map VName TypeBinding -> Maybe TypeBinding
forall a b. (a -> b) -> a -> b
$ Env -> Map VName TypeBinding
envTypeTable Env
scope of
      Maybe TypeBinding
Nothing -> String -> TypeM ([TypeParam], StructRetType, Liftedness)
forall a. HasCallStack => String -> a
error (String -> TypeM ([TypeParam], StructRetType, Liftedness))
-> String -> TypeM ([TypeParam], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$ String
"lookupType: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> QualName VName -> String
forall a. Show a => a -> String
show QualName VName
qn
      Just (TypeAbbr Liftedness
l [TypeParam]
ps (RetType [VName]
dims TypeBase Size NoUniqueness
def)) ->
        ([TypeParam], StructRetType, Liftedness)
-> TypeM ([TypeParam], StructRetType, Liftedness)
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TypeParam]
ps, [VName] -> TypeBase Size NoUniqueness -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase Size NoUniqueness -> StructRetType)
-> TypeBase Size NoUniqueness -> StructRetType
forall a b. (a -> b) -> a -> b
$ Env
-> [VName]
-> [VName]
-> TypeBase Size NoUniqueness
-> TypeBase Size NoUniqueness
forall as.
Env -> [VName] -> [VName] -> TypeBase Size as -> TypeBase Size as
qualifyTypeVars Env
outer_env [VName]
forall a. Monoid a => a
mempty (QualName VName -> [VName]
forall vn. QualName vn -> [vn]
qualQuals QualName VName
qn) TypeBase Size NoUniqueness
def, Liftedness
l)

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

lookupQualNameEnv :: QualName VName -> TypeM Env
lookupQualNameEnv :: QualName VName -> TypeM Env
lookupQualNameEnv qn :: QualName VName
qn@(QualName [VName]
quals VName
_) = do
  Env
env <- TypeM Env
askEnv
  Env -> [VName] -> TypeM Env
forall {f :: * -> *}. Applicative f => Env -> [VName] -> f Env
descend Env
env [VName]
quals
  where
    descend :: Env -> [VName] -> f Env
descend Env
scope [] = Env -> f Env
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
scope
    descend Env
scope (VName
q : [VName]
qs)
      | Just (ModEnv Env
q_scope) <- VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q (Map VName Mod -> Maybe Mod) -> Map VName Mod -> Maybe Mod
forall a b. (a -> b) -> a -> b
$ Env -> Map VName Mod
envModTable Env
scope =
          Env -> [VName] -> f Env
descend Env
q_scope [VName]
qs
      | Bool
otherwise =
          String -> f Env
forall a. HasCallStack => String -> a
error (String -> f Env) -> String -> f Env
forall a b. (a -> b) -> a -> b
$ String
"lookupQualNameEnv: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ QualName VName -> String
forall a. Show a => a -> String
show QualName VName
qn

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

-- | Elaborate the given name in the given namespace at the given
-- location, producing the corresponding unique 'VName'.
checkName :: Namespace -> Name -> SrcLoc -> TypeM VName
checkName :: Namespace -> Name -> SrcLoc -> TypeM VName
checkName Namespace
space Name
name SrcLoc
loc = QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf (QualName VName -> VName) -> TypeM (QualName VName) -> TypeM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Namespace -> QualName Name -> SrcLoc -> TypeM (QualName VName)
checkQualName Namespace
space (Name -> QualName Name
forall v. v -> QualName v
qualName Name
name) SrcLoc
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 :: QualName VName -> TypeM Bool
isKnownType QualName VName
qn = do
  Env
env <- TypeM Env
askEnv
  Env -> [VName] -> VName -> TypeM Bool
forall {f :: * -> *}.
Applicative f =>
Env -> [VName] -> VName -> f Bool
descend Env
env (QualName VName -> [VName]
forall vn. QualName vn -> [vn]
qualQuals QualName VName
qn) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
  where
    descend :: Env -> [VName] -> VName -> f Bool
descend Env
env [] VName
v
      | Just QualName VName
v' <- (Namespace, Name) -> NameMap -> Maybe (QualName VName)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
Type, VName -> Name
baseName VName
v) (NameMap -> Maybe (QualName VName))
-> NameMap -> Maybe (QualName VName)
forall a b. (a -> b) -> a -> b
$ Env -> NameMap
envNameMap Env
env =
          Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> f Bool) -> Bool -> f Bool
forall a b. (a -> b) -> a -> b
$ VName -> Map VName TypeBinding -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v') (Map VName TypeBinding -> Bool) -> Map VName TypeBinding -> Bool
forall a b. (a -> b) -> a -> b
$ Env -> Map VName TypeBinding
envTypeTable Env
env
    descend Env
env (VName
q : [VName]
qs) VName
v
      | Just QualName VName
q' <- (Namespace, Name) -> NameMap -> Maybe (QualName VName)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
Term, VName -> Name
baseName VName
q) (NameMap -> Maybe (QualName VName))
-> NameMap -> Maybe (QualName VName)
forall a b. (a -> b) -> a -> b
$ Env -> NameMap
envNameMap Env
env,
        Just (ModEnv Env
env') <- VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
q') (Map VName Mod -> Maybe Mod) -> Map VName Mod -> Maybe Mod
forall a b. (a -> b) -> a -> b
$ Env -> Map VName Mod
envModTable Env
env =
          Env -> [VName] -> VName -> f Bool
descend Env
env' [VName]
qs VName
v
    descend Env
_ [VName]
_ VName
_ = Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

lookupMod :: SrcLoc -> QualName Name -> TypeM (QualName VName, Mod)
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 VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName Mod -> Maybe Mod) -> Map VName Mod -> Maybe Mod
forall a b. (a -> b) -> a -> b
$ Env -> Map VName Mod
envModTable Env
scope of
    Maybe Mod
Nothing -> Namespace -> QualName Name -> SrcLoc -> TypeM (QualName VName, Mod)
forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
Term QualName Name
qn SrcLoc
loc
    Just Mod
m -> (QualName VName, Mod) -> TypeM (QualName VName, Mod)
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName
qn', Mod
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 :: forall as.
Env -> [VName] -> [VName] -> TypeBase Size as -> TypeBase Size as
qualifyTypeVars Env
outer_env [VName]
orig_except [VName]
ref_qs = Set VName -> TypeBase Size as -> TypeBase Size as
forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType ([VName] -> Set VName
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
u Shape Size
shape ScalarTypeBase Size NoUniqueness
et) =
      as
-> Shape Size
-> ScalarTypeBase Size NoUniqueness
-> TypeBase Size as
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array as
u ((Size -> Size) -> Shape Size -> Shape Size
forall a b. (a -> b) -> Shape a -> Shape b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Set VName -> Size -> Size
forall {t :: * -> *}. Foldable t => t VName -> Size -> Size
onDim Set VName
except) Shape Size
shape) (Set VName
-> ScalarTypeBase Size NoUniqueness
-> ScalarTypeBase Size NoUniqueness
forall {u}.
Set VName -> ScalarTypeBase Size u -> ScalarTypeBase Size u
onScalar Set VName
except ScalarTypeBase Size NoUniqueness
et)
    onType Set VName
except (Scalar ScalarTypeBase Size as
t) =
      ScalarTypeBase Size as -> TypeBase Size as
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size as -> TypeBase Size as)
-> ScalarTypeBase Size as -> TypeBase Size as
forall a b. (a -> b) -> a -> b
$ Set VName -> ScalarTypeBase Size as -> ScalarTypeBase Size as
forall {u}.
Set VName -> ScalarTypeBase Size u -> ScalarTypeBase Size u
onScalar Set VName
except ScalarTypeBase Size as
t

    onScalar :: Set VName -> ScalarTypeBase Size u -> ScalarTypeBase Size u
onScalar Set VName
_ (Prim PrimType
t) = PrimType -> ScalarTypeBase Size u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t
    onScalar Set VName
except (TypeVar u
u QualName VName
qn [TypeArg Size]
targs) =
      u -> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size u
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar u
u (Set VName -> QualName VName -> QualName VName
forall {t :: * -> *}.
Foldable t =>
t VName -> QualName VName -> QualName VName
qual Set VName
except QualName VName
qn) ((TypeArg Size -> TypeArg Size) -> [TypeArg Size] -> [TypeArg Size]
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 u)
m) =
      Map Name (TypeBase Size u) -> ScalarTypeBase Size u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase Size u) -> ScalarTypeBase Size u)
-> Map Name (TypeBase Size u) -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ (TypeBase Size u -> TypeBase Size u)
-> Map Name (TypeBase Size u) -> Map Name (TypeBase Size u)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Set VName -> TypeBase Size u -> TypeBase Size u
forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except) Map Name (TypeBase Size u)
m
    onScalar Set VName
except (Sum Map Name [TypeBase Size u]
m) =
      Map Name [TypeBase Size u] -> ScalarTypeBase Size u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase Size u] -> ScalarTypeBase Size u)
-> Map Name [TypeBase Size u] -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ ([TypeBase Size u] -> [TypeBase Size u])
-> Map Name [TypeBase Size u] -> Map Name [TypeBase Size u]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase Size u -> TypeBase Size u)
-> [TypeBase Size u] -> [TypeBase Size u]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase Size u -> TypeBase Size u)
 -> [TypeBase Size u] -> [TypeBase Size u])
-> (TypeBase Size u -> TypeBase Size u)
-> [TypeBase Size u]
-> [TypeBase Size u]
forall a b. (a -> b) -> a -> b
$ Set VName -> TypeBase Size u -> TypeBase Size u
forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except) Map Name [TypeBase Size u]
m
    onScalar Set VName
except (Arrow u
as PName
p Diet
d TypeBase Size NoUniqueness
t1 (RetType [VName]
dims TypeBase Size Uniqueness
t2)) =
      u
-> PName
-> Diet
-> TypeBase Size NoUniqueness
-> RetTypeBase Size Uniqueness
-> ScalarTypeBase Size u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
as PName
p Diet
d (Set VName
-> TypeBase Size NoUniqueness -> TypeBase Size NoUniqueness
forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except' TypeBase Size NoUniqueness
t1) (RetTypeBase Size Uniqueness -> ScalarTypeBase Size u)
-> RetTypeBase Size Uniqueness -> ScalarTypeBase Size u
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Uniqueness -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (Set VName -> TypeBase Size Uniqueness -> TypeBase Size Uniqueness
forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except' TypeBase Size Uniqueness
t2)
      where
        except' :: Set VName
except' = case PName
p of
          Named VName
p' -> VName -> Set VName -> Set VName
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) =
      Size -> TypeArg Size
forall dim. dim -> TypeArg dim
TypeArgDim (Size -> TypeArg Size) -> Size -> TypeArg Size
forall a b. (a -> b) -> a -> b
$ Set VName -> Size -> Size
forall {t :: * -> *}. Foldable t => t VName -> Size -> Size
onDim Set VName
except Size
d
    onTypeArg Set VName
except (TypeArgType TypeBase Size NoUniqueness
t) =
      TypeBase Size NoUniqueness -> TypeArg Size
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (TypeBase Size NoUniqueness -> TypeArg Size)
-> TypeBase Size NoUniqueness -> TypeArg Size
forall a b. (a -> b) -> a -> b
$ Set VName
-> TypeBase Size NoUniqueness -> TypeBase Size NoUniqueness
forall as. Set VName -> TypeBase Size as -> TypeBase Size as
onType Set VName
except TypeBase Size NoUniqueness
t

    onDim :: t VName -> Size -> Size
onDim t VName
except Size
e = Identity Size -> Size
forall a. Identity a -> a
runIdentity (Identity Size -> Size) -> Identity Size -> Size
forall a b. (a -> b) -> a -> b
$ t VName -> Size -> Identity Size
forall {f :: * -> *} {t :: * -> *}.
(Foldable t, Monad f) =>
t VName -> Size -> f Size
onDimM t VName
except Size
e
    onDimM :: t VName -> Size -> f Size
onDimM t VName
except (Var QualName VName
qn Info (TypeBase Size NoUniqueness)
typ SrcLoc
loc) = Size -> f Size
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Size -> f Size) -> Size -> f Size
forall a b. (a -> b) -> a -> b
$ QualName VName
-> Info (TypeBase Size NoUniqueness) -> SrcLoc -> Size
forall (f :: * -> *) vn.
QualName vn
-> f (TypeBase Size NoUniqueness) -> SrcLoc -> ExpBase f vn
Var (t VName -> QualName VName -> QualName VName
forall {t :: * -> *}.
Foldable t =>
t VName -> QualName VName -> QualName VName
qual t VName
except QualName VName
qn) Info (TypeBase Size NoUniqueness)
typ SrcLoc
loc
    onDimM t VName
except Size
e = ASTMapper f -> Size -> f Size
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Size -> m Size
astMap (ASTMapper f
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = onDimM except}) Size
e

    qual :: t VName -> QualName VName -> QualName VName
qual t VName
except (QualName [VName]
orig_qs VName
name)
      | VName
name VName -> t VName -> Bool
forall a. Eq a => a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
except Bool -> Bool -> Bool
|| [VName] -> VName -> Env -> Bool
reachable [VName]
orig_qs VName
name Env
outer_env =
          [VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name
      | Bool
otherwise =
          [VName] -> [VName] -> QualName VName -> QualName VName
prependAsNecessary [] [VName]
ref_qs (QualName VName -> QualName VName)
-> QualName VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ [VName] -> VName -> QualName VName
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 [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
orig_qs) VName
name Env
outer_env = [VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName ([VName]
qs [VName] -> [VName] -> [VName]
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 [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName
q]) [VName]
rem_qs' ([VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name)
          [] -> [VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName [VName]
orig_qs VName
name

    reachable :: [VName] -> VName -> Env -> Bool
reachable [] VName
name Env
env =
      VName
name VName -> Map VName BoundV -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Env -> Map VName BoundV
envVtable Env
env
        Bool -> Bool -> Bool
|| Maybe TypeBinding -> Bool
forall a. Maybe a -> Bool
isJust ((TypeBinding -> Bool) -> [TypeBinding] -> Maybe TypeBinding
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find TypeBinding -> Bool
matches ([TypeBinding] -> Maybe TypeBinding)
-> [TypeBinding] -> Maybe TypeBinding
forall a b. (a -> b) -> a -> b
$ Map VName TypeBinding -> [TypeBinding]
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 NoUniqueness
_ (QualName [VName]
x_qs VName
name') [TypeArg Size]
_)))) =
          [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
x_qs Bool -> Bool -> Bool
&& VName
name VName -> VName -> Bool
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') <- VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q (Map VName Mod -> Maybe Mod) -> Map VName Mod -> Maybe Mod
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 = (TypeError -> TypeM a)
-> (a -> TypeM a) -> Either TypeError a -> TypeM a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either TypeError -> TypeM a
forall a. TypeError -> TypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError a -> TypeM a
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

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

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

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

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

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

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

--- Name handling

-- | The 'NameMap' corresponding to the intrinsics module.
intrinsicsNameMap :: NameMap
intrinsicsNameMap :: NameMap
intrinsicsNameMap = [((Namespace, Name), QualName VName)] -> NameMap
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([((Namespace, Name), QualName VName)] -> NameMap)
-> [((Namespace, Name), QualName VName)] -> NameMap
forall a b. (a -> b) -> a -> b
$ ((VName, Intrinsic) -> ((Namespace, Name), QualName VName))
-> [(VName, Intrinsic)] -> [((Namespace, Name), QualName VName)]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Intrinsic) -> ((Namespace, Name), QualName VName)
mapping ([(VName, Intrinsic)] -> [((Namespace, Name), QualName VName)])
-> [(VName, Intrinsic)] -> [((Namespace, Name), QualName VName)]
forall a b. (a -> b) -> a -> b
$ Map VName Intrinsic -> [(VName, Intrinsic)]
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), [VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName [] VName
v)
    mapping (VName
v, Intrinsic
_) = ((Namespace
Term, VName -> Name
baseName VName
v), [VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName [] VName
v)

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

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