{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TupleSections #-}

-- | Main monad in which the type checker runs, as well as ancillary
-- data definitions.
module Language.Futhark.TypeChecker.Monad
  ( TypeM,
    runTypeM,
    askEnv,
    askImportName,
    bindSpaced,
    qualifyTypeVars,
    lookupMTy,
    lookupImport,
    localEnv,
    TypeError (..),
    unappliedFunctor,
    unknownVariable,
    unknownType,
    underscoreUse,
    Notes,
    aNote,
    MonadTypeChecker (..),
    checkName,
    badOnLeft,
    module Language.Futhark.Warnings,
    Env (..),
    TySet,
    FunSig (..),
    ImportTable,
    NameMap,
    BoundV (..),
    Mod (..),
    TypeBinding (..),
    MTy (..),
    anySignedType,
    anyUnsignedType,
    anyFloatType,
    anyNumberType,
    anyPrimType,
    Namespace (..),
    intrinsicsNameMap,
    topLevelNameMap,
  )
where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Either
import Data.List (find, isPrefixOf)
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import Futhark.FreshNames hiding (newName)
import qualified Futhark.FreshNames
import Futhark.Util.Console
import Futhark.Util.Pretty hiding (space)
import Language.Futhark
import Language.Futhark.Semantic
import Language.Futhark.Warnings
import Prelude hiding (mapM, mod)

-- | A note with extra information regarding a type error.
newtype Note = Note Doc

-- | A collection of 'Note's.
newtype Notes = Notes [Note]
  deriving (b -> Notes -> Notes
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
stimes :: 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
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
mconcat :: [Notes] -> Notes
$cmconcat :: [Notes] -> Notes
mappend :: Notes -> Notes -> Notes
$cmappend :: Notes -> Notes -> Notes
mempty :: Notes
$cmempty :: Notes
$cp1Monoid :: Semigroup Notes
Monoid)

instance Pretty Note where
  ppr :: Note -> Doc
ppr (Note Doc
msg) = Doc
"Note:" Doc -> Doc -> Doc
<+> Doc -> Doc
align Doc
msg

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

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

-- | Information about an error during type checking.
data TypeError = TypeError SrcLoc Notes Doc

instance Pretty TypeError where
  ppr :: TypeError -> Doc
ppr (TypeError SrcLoc
loc Notes
notes Doc
msg) =
    String -> Doc
text (String -> String
inRed (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ String
"Error at " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
":")
      Doc -> Doc -> Doc
</> Doc
msg Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Notes -> Doc
forall a. Pretty a => a -> Doc
ppr Notes
notes

-- | An unexpected functor appeared!
unappliedFunctor :: MonadTypeChecker m => SrcLoc -> m a
unappliedFunctor :: SrcLoc -> m a
unappliedFunctor SrcLoc
loc =
  SrcLoc -> 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 :: Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
name SrcLoc
loc =
  SrcLoc -> 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
<+> Namespace -> Doc
forall a. Pretty a => a -> Doc
ppr Namespace
space Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr QualName Name
name)

-- | An unknown type was referenced.
unknownType :: MonadTypeChecker m => SrcLoc -> QualName Name -> m a
unknownType :: SrcLoc -> QualName Name -> m a
unknownType SrcLoc
loc QualName Name
name =
  SrcLoc -> 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 type" Doc -> Doc -> Doc
<+> QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr QualName Name
name Doc -> Doc -> Doc
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 :: SrcLoc -> QualName Name -> m a
underscoreUse SrcLoc
loc QualName Name
name =
  SrcLoc -> 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
<+> Doc -> Doc
pquote (QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr 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 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
  }

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

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

instance MonadError TypeError TypeM where
  throwError :: 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 e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Warnings
ws, TypeError
e)

  catchError :: 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 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 :: 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 -> Context
Context Env
env ImportTable
imports ImportName
fpath
      s :: TypeState
s = VNameSource -> Warnings -> TypeState
TypeState VNameSource
src Warnings
forall a. Monoid a => a
mempty
  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 VNameSource
src' Warnings
ws) -> (Warnings
ws, (a, VNameSource) -> Either TypeError (a, VNameSource)
forall a b. b -> Either a b
Right (a
x, VNameSource
src'))

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

-- | 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 (m :: * -> *) a. Monad m => a -> m a
return (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
envSigTable 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 (FilePath, Env)
lookupImport :: SrcLoc -> String -> TypeM (String, 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 :: String
canonical_import = ImportName -> String
includeToString (ImportName -> String) -> ImportName -> String
forall a b. (a -> b) -> a -> b
$ ImportName -> String -> SrcLoc -> ImportName
mkImportFrom ImportName
my_path String
file SrcLoc
loc
  case String -> ImportTable -> Maybe Env
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
canonical_import ImportTable
imports of
    Maybe Env
Nothing ->
      SrcLoc -> Notes -> Doc -> TypeM (String, Env)
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 (String, Env)) -> Doc -> TypeM (String, Env)
forall a b. (a -> b) -> a -> b
$
        Doc
"Unknown import" Doc -> Doc -> Doc
<+> Doc -> Doc
dquotes (String -> Doc
text String
canonical_import)
          Doc -> Doc -> Doc
</> Doc
"Known:" Doc -> Doc -> Doc
<+> [Doc] -> Doc
commasep ((String -> Doc) -> [String] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map String -> Doc
text (ImportTable -> [String]
forall k a. Map k a -> [k]
M.keys ImportTable
imports))
    Just Env
scope -> (String, Env) -> TypeM (String, Env)
forall (m :: * -> *) a. Monad m => a -> m a
return (String
canonical_import, Env
scope)

-- | Evaluate a 'TypeM' computation within an extended (/not/
-- replaced) environment.
localEnv :: Env -> TypeM a -> TypeM a
localEnv :: Env -> TypeM a -> TypeM a
localEnv Env
env = (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
contextEnv = Env
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 ()

  newName :: VName -> m VName
  newID :: 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], StructType, Liftedness)
  lookupMod :: SrcLoc -> QualName Name -> m (QualName VName, Mod)
  lookupVar :: SrcLoc -> QualName Name -> m (QualName VName, PatternType)

  checkNamedDim :: SrcLoc -> QualName Name -> m (QualName VName)
  checkNamedDim SrcLoc
loc QualName Name
v = do
    (QualName VName
v', PatternType
t) <- SrcLoc -> QualName Name -> m (QualName VName, PatternType)
forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m (QualName VName, PatternType)
lookupVar SrcLoc
loc QualName Name
v
    case PatternType
t of
      Scalar (Prim (Signed IntType
Int64)) -> QualName VName -> m (QualName VName)
forall (m :: * -> *) a. Monad m => a -> m a
return QualName VName
v'
      PatternType
_ ->
        SrcLoc -> Notes -> Doc -> m (QualName VName)
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 (QualName VName)) -> Doc -> m (QualName VName)
forall a b. (a -> b) -> a -> b
$
          Doc
"Dimension declaration" Doc -> Doc -> Doc
<+> QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr QualName Name
v Doc -> Doc -> Doc
<+> Doc
"should be of type i64."

  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 :: Namespace -> Name -> SrcLoc -> m VName
checkName Namespace
space Name
name SrcLoc
loc = QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf (QualName VName -> VName) -> m (QualName VName) -> m VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Namespace -> QualName Name -> SrcLoc -> m (QualName VName)
forall (m :: * -> *).
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m (QualName VName)
checkQualName Namespace
space (Name -> QualName Name
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 :: [(Namespace, Name)] -> m a -> m a
bindSpaced [(Namespace, Name)]
names m a
body = do
  [VName]
names' <- ((Namespace, Name) -> m VName) -> [(Namespace, Name)] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name -> m VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID (Name -> m VName)
-> ((Namespace, Name) -> Name) -> (Namespace, Name) -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Namespace, Name) -> Name
forall a b. (a, b) -> b
snd) [(Namespace, Name)]
names
  let mapping :: Map (Namespace, Name) (QualName VName)
mapping = [((Namespace, Name), QualName VName)]
-> Map (Namespace, Name) (QualName VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Namespace, Name)]
-> [QualName VName] -> [((Namespace, Name), QualName VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Namespace, Name)]
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')
  Map (Namespace, Name) (QualName VName) -> m a -> m a
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 :: loc -> Doc -> TypeM ()
warn loc
loc Doc
problem =
    (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 :: Warnings
stateWarnings = TypeState -> Warnings
stateWarnings TypeState
s Warnings -> Warnings -> Warnings
forall a. Semigroup a => a -> a -> a
<> SrcLoc -> Doc -> Warnings
singleWarning (loc -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf 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 :: VNameSource
stateNameSource = VNameSource
src'}
    VName -> TypeM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
v'

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

  bindNameMap :: Map (Namespace, Name) (QualName VName) -> TypeM a -> TypeM a
bindNameMap Map (Namespace, Name) (QualName VName)
m = (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
contextEnv = Env
env {envNameMap :: Map (Namespace, Name) (QualName VName)
envNameMap = Map (Namespace, Name) (QualName VName)
m Map (Namespace, Name) (QualName VName)
-> Map (Namespace, Name) (QualName VName)
-> Map (Namespace, Name) (QualName VName)
forall a. Semigroup a => a -> a -> a
<> Env -> Map (Namespace, Name) (QualName VName)
envNameMap Env
env}}

  bindVal :: VName -> BoundV -> TypeM a -> TypeM a
bindVal VName
v BoundV
t = (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 :: Env
contextEnv =
          (Context -> Env
contextEnv Context
ctx)
            { envVtable :: Map VName BoundV
envVtable = VName -> BoundV -> Map VName BoundV -> Map VName BoundV
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v BoundV
t (Map VName BoundV -> Map VName BoundV)
-> Map VName BoundV -> Map VName BoundV
forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable (Env -> Map VName BoundV) -> Env -> Map VName BoundV
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 = (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

  lookupType :: SrcLoc
-> QualName Name
-> TypeM (QualName VName, [TypeParam], StructType, 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 VName -> Map VName TypeBinding -> Maybe TypeBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (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 -> SrcLoc
-> QualName Name
-> TypeM (QualName VName, [TypeParam], StructType, Liftedness)
forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
unknownType SrcLoc
loc QualName Name
qn
      Just (TypeAbbr Liftedness
l [TypeParam]
ps StructType
def) -> (QualName VName, [TypeParam], StructType, Liftedness)
-> TypeM (QualName VName, [TypeParam], StructType, Liftedness)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
qn', [TypeParam]
ps, Env -> [VName] -> [VName] -> StructType -> StructType
forall as.
Env
-> [VName]
-> [VName]
-> TypeBase (DimDecl VName) as
-> TypeBase (DimDecl VName) as
qualifyTypeVars Env
outer_env [VName]
forall a. Monoid a => a
mempty [VName]
qs StructType
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 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 (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
qn', Mod
m)

  lookupVar :: SrcLoc -> QualName Name -> TypeM (QualName VName, PatternType)
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 VName -> Map VName BoundV -> Maybe BoundV
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName BoundV -> Maybe BoundV)
-> Map VName BoundV -> Maybe BoundV
forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable Env
env of
      Maybe BoundV
Nothing -> Namespace
-> QualName Name -> SrcLoc -> TypeM (QualName VName, PatternType)
forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
Term QualName Name
qn SrcLoc
loc
      Just (BoundV [TypeParam]
_ StructType
t)
        | String
"_" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` VName -> String
baseString VName
name -> SrcLoc -> QualName Name -> TypeM (QualName VName, PatternType)
forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
underscoreUse SrcLoc
loc QualName Name
qn
        | Bool
otherwise ->
          case StructType -> Either ([(PName, StructType)], StructType) StructType
forall dim as.
TypeBase dim as
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
getType StructType
t of
            Left {} ->
              SrcLoc -> Notes -> Doc -> TypeM (QualName VName, PatternType)
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 (QualName VName, PatternType))
-> Doc -> TypeM (QualName VName, PatternType)
forall a b. (a -> b) -> a -> b
$
                Doc
"Attempt to use function" Doc -> Doc -> Doc
<+> VName -> Doc
forall v. IsName v => v -> Doc
pprName VName
name Doc -> Doc -> Doc
<+> Doc
"as value."
            Right StructType
t' ->
              (QualName VName, PatternType)
-> TypeM (QualName VName, PatternType)
forall (m :: * -> *) a. Monad m => a -> m a
return
                ( QualName VName
qn',
                  StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct (StructType -> PatternType) -> StructType -> PatternType
forall a b. (a -> b) -> a -> b
$
                    Env -> [VName] -> [VName] -> StructType -> StructType
forall as.
Env
-> [VName]
-> [VName]
-> TypeBase (DimDecl VName) as
-> TypeBase (DimDecl VName) as
qualifyTypeVars Env
outer_env [VName]
forall a. Monoid a => a
mempty [VName]
qs StructType
t'
                )

  typeError :: loc -> Notes -> Doc -> TypeM a
typeError loc
loc Notes
notes Doc
s = 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
$ SrcLoc -> Notes -> Doc -> TypeError
TypeError (loc -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf loc
loc) Notes
notes Doc
s

-- | Extract from a type either a function type comprising a list of
-- parameter types and a return type, or a first-order type.
getType ::
  TypeBase dim as ->
  Either
    ([(PName, TypeBase dim as)], TypeBase dim as)
    (TypeBase dim as)
getType :: TypeBase dim as
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
getType (Scalar (Arrow as
_ PName
v TypeBase dim as
t1 TypeBase dim as
t2)) =
  case TypeBase dim as
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
forall dim as.
TypeBase dim as
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
getType TypeBase dim as
t2 of
    Left ([(PName, TypeBase dim as)]
ps, TypeBase dim as
r) -> ([(PName, TypeBase dim as)], TypeBase dim as)
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
forall a b. a -> Either a b
Left ((PName
v, TypeBase dim as
t1) (PName, TypeBase dim as)
-> [(PName, TypeBase dim as)] -> [(PName, TypeBase dim as)]
forall a. a -> [a] -> [a]
: [(PName, TypeBase dim as)]
ps, TypeBase dim as
r)
    Right TypeBase dim as
_ -> ([(PName, TypeBase dim as)], TypeBase dim as)
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
forall a b. a -> Either a b
Left ([(PName
v, TypeBase dim as
t1)], TypeBase dim as
t2)
getType TypeBase dim as
t = TypeBase dim as
-> Either
     ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as)
forall a b. b -> Either a b
Right 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
  Env -> [Name] -> TypeM (Env, QualName VName)
forall (m :: * -> *).
MonadTypeChecker m =>
Env -> [Name] -> m (Env, QualName VName)
descend Env
env [Name]
quals
  where
    descend :: Env -> [Name] -> m (Env, QualName VName)
descend Env
scope []
      | Just QualName VName
name' <- (Namespace, Name)
-> Map (Namespace, Name) (QualName VName) -> Maybe (QualName VName)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
space, Name
name) (Map (Namespace, Name) (QualName VName) -> Maybe (QualName VName))
-> Map (Namespace, Name) (QualName VName) -> Maybe (QualName VName)
forall a b. (a -> b) -> a -> b
$ Env -> Map (Namespace, Name) (QualName VName)
envNameMap Env
scope =
        (Env, QualName VName) -> m (Env, QualName VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (Env
scope, QualName VName
name')
      | Bool
otherwise =
        Namespace -> QualName Name -> SrcLoc -> m (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)
-> Map (Namespace, Name) (QualName VName) -> Maybe (QualName VName)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
Term, Name
q) (Map (Namespace, Name) (QualName VName) -> Maybe (QualName VName))
-> Map (Namespace, Name) (QualName VName) -> Maybe (QualName VName)
forall a b. (a -> b) -> a -> b
$ Env -> Map (Namespace, Name) (QualName VName)
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 =
        case Mod
res of
          ModEnv Env
q_scope -> do
            (Env
scope', QualName [VName]
qs' VName
name') <- Env -> [Name] -> m (Env, QualName VName)
descend Env
q_scope [Name]
qs
            (Env, QualName VName) -> m (Env, QualName VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (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 -> m (Env, QualName VName)
forall (m :: * -> *) a. MonadTypeChecker m => SrcLoc -> m a
unappliedFunctor SrcLoc
loc
      | Bool
otherwise =
        Namespace -> QualName Name -> SrcLoc -> m (Env, QualName VName)
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 (DimDecl VName) as ->
  TypeBase (DimDecl VName) as
qualifyTypeVars :: Env
-> [VName]
-> [VName]
-> TypeBase (DimDecl VName) as
-> TypeBase (DimDecl VName) as
qualifyTypeVars Env
outer_env [VName]
orig_except [VName]
ref_qs = Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall as.
Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
orig_except)
  where
    onType ::
      S.Set VName ->
      TypeBase (DimDecl VName) as ->
      TypeBase (DimDecl VName) as
    onType :: Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType Set VName
except (Array as
as Uniqueness
u ScalarTypeBase (DimDecl VName) ()
et ShapeDecl (DimDecl VName)
shape) =
      as
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) as
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array as
as Uniqueness
u (Set VName
-> ScalarTypeBase (DimDecl VName) ()
-> ScalarTypeBase (DimDecl VName) ()
forall as.
Set VName
-> ScalarTypeBase (DimDecl VName) as
-> ScalarTypeBase (DimDecl VName) as
onScalar Set VName
except ScalarTypeBase (DimDecl VName) ()
et) ((DimDecl VName -> DimDecl VName)
-> ShapeDecl (DimDecl VName) -> ShapeDecl (DimDecl VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Set VName -> DimDecl VName -> DimDecl VName
forall (t :: * -> *).
Foldable t =>
t VName -> DimDecl VName -> DimDecl VName
onDim Set VName
except) ShapeDecl (DimDecl VName)
shape)
    onType Set VName
except (Scalar ScalarTypeBase (DimDecl VName) as
t) =
      ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ Set VName
-> ScalarTypeBase (DimDecl VName) as
-> ScalarTypeBase (DimDecl VName) as
forall as.
Set VName
-> ScalarTypeBase (DimDecl VName) as
-> ScalarTypeBase (DimDecl VName) as
onScalar Set VName
except ScalarTypeBase (DimDecl VName) as
t

    onScalar :: Set VName
-> ScalarTypeBase (DimDecl VName) as
-> ScalarTypeBase (DimDecl VName) as
onScalar Set VName
_ (Prim PrimType
t) = PrimType -> ScalarTypeBase (DimDecl VName) as
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    onScalar Set VName
except (TypeVar as
as Uniqueness
u TypeName
tn [TypeArg (DimDecl VName)]
targs) =
      as
-> Uniqueness
-> TypeName
-> [TypeArg (DimDecl VName)]
-> ScalarTypeBase (DimDecl VName) as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar as
as Uniqueness
u TypeName
tn' ([TypeArg (DimDecl VName)] -> ScalarTypeBase (DimDecl VName) as)
-> [TypeArg (DimDecl VName)] -> ScalarTypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ (TypeArg (DimDecl VName) -> TypeArg (DimDecl VName))
-> [TypeArg (DimDecl VName)] -> [TypeArg (DimDecl VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Set VName -> TypeArg (DimDecl VName) -> TypeArg (DimDecl VName)
onTypeArg Set VName
except) [TypeArg (DimDecl VName)]
targs
      where
        tn' :: TypeName
tn' = QualName VName -> TypeName
typeNameFromQualName (QualName VName -> TypeName) -> QualName VName -> TypeName
forall a b. (a -> b) -> a -> b
$ Set VName -> QualName VName -> QualName VName
forall (t :: * -> *).
Foldable t =>
t VName -> QualName VName -> QualName VName
qual Set VName
except (QualName VName -> QualName VName)
-> QualName VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ TypeName -> QualName VName
qualNameFromTypeName TypeName
tn
    onScalar Set VName
except (Record Map Name (TypeBase (DimDecl VName) as)
m) =
      Map Name (TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase (DimDecl VName) as)
 -> ScalarTypeBase (DimDecl VName) as)
-> Map Name (TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ (TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> Map Name (TypeBase (DimDecl VName) as)
-> Map Name (TypeBase (DimDecl VName) as)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall as.
Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType Set VName
except) Map Name (TypeBase (DimDecl VName) as)
m
    onScalar Set VName
except (Sum Map Name [TypeBase (DimDecl VName) as]
m) =
      Map Name [TypeBase (DimDecl VName) as]
-> ScalarTypeBase (DimDecl VName) as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase (DimDecl VName) as]
 -> ScalarTypeBase (DimDecl VName) as)
-> Map Name [TypeBase (DimDecl VName) as]
-> ScalarTypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ ([TypeBase (DimDecl VName) as] -> [TypeBase (DimDecl VName) as])
-> Map Name [TypeBase (DimDecl VName) as]
-> Map Name [TypeBase (DimDecl VName) as]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> [TypeBase (DimDecl VName) as] -> [TypeBase (DimDecl VName) as]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
 -> [TypeBase (DimDecl VName) as] -> [TypeBase (DimDecl VName) as])
-> (TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> [TypeBase (DimDecl VName) as]
-> [TypeBase (DimDecl VName) as]
forall a b. (a -> b) -> a -> b
$ Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall as.
Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType Set VName
except) Map Name [TypeBase (DimDecl VName) as]
m
    onScalar Set VName
except (Arrow as
as PName
p TypeBase (DimDecl VName) as
t1 TypeBase (DimDecl VName) as
t2) =
      as
-> PName
-> TypeBase (DimDecl VName) as
-> TypeBase (DimDecl VName) as
-> ScalarTypeBase (DimDecl VName) as
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow as
as PName
p (Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall as.
Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType Set VName
except' TypeBase (DimDecl VName) as
t1) (Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall as.
Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType Set VName
except' TypeBase (DimDecl VName) as
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 (DimDecl VName) -> TypeArg (DimDecl VName)
onTypeArg Set VName
except (TypeArgDim DimDecl VName
d SrcLoc
loc) =
      DimDecl VName -> SrcLoc -> TypeArg (DimDecl VName)
forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim (Set VName -> DimDecl VName -> DimDecl VName
forall (t :: * -> *).
Foldable t =>
t VName -> DimDecl VName -> DimDecl VName
onDim Set VName
except DimDecl VName
d) SrcLoc
loc
    onTypeArg Set VName
except (TypeArgType StructType
t SrcLoc
loc) =
      StructType -> SrcLoc -> TypeArg (DimDecl VName)
forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (Set VName -> StructType -> StructType
forall as.
Set VName
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
onType Set VName
except StructType
t) SrcLoc
loc

    onDim :: t VName -> DimDecl VName -> DimDecl VName
onDim t VName
except (NamedDim QualName VName
qn) = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ t VName -> QualName VName -> QualName VName
forall (t :: * -> *).
Foldable t =>
t VName -> QualName VName -> QualName VName
qual t VName
except QualName VName
qn
    onDim t VName
_ DimDecl VName
d = DimDecl VName
d

    qual :: t VName -> QualName VName -> QualName VName
qual t VName
except (QualName [VName]
orig_qs VName
name)
      | VName
name VName -> t VName -> 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]
_ (Scalar (TypeVar ()
_ Uniqueness
_ (TypeName [VName]
x_qs VName
name') [TypeArg (DimDecl VName)]
_))) =
          [VName] -> 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 :: 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 e (m :: * -> *) a. MonadError e m => e -> m a
throwError a -> TypeM a
forall (m :: * -> *) a. Monad m => a -> m a
return

-- | 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 :: Map (Namespace, Name) (QualName VName)
intrinsicsNameMap = [((Namespace, Name), QualName VName)]
-> Map (Namespace, Name) (QualName VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([((Namespace, Name), QualName VName)]
 -> Map (Namespace, Name) (QualName VName))
-> [((Namespace, Name), QualName VName)]
-> Map (Namespace, Name) (QualName VName)
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 :: Map (Namespace, Name) (QualName VName)
topLevelNameMap = ((Namespace, Name) -> QualName VName -> Bool)
-> Map (Namespace, Name) (QualName VName)
-> Map (Namespace, Name) (QualName VName)
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\(Namespace, Name)
k QualName VName
_ -> (Namespace, Name) -> Bool
atTopLevel (Namespace, Name)
k) Map (Namespace, Name) (QualName VName)
intrinsicsNameMap
  where
    atTopLevel :: (Namespace, Name) -> Bool
    atTopLevel :: (Namespace, Name) -> Bool
atTopLevel (Namespace
Type, Name
_) = Bool
True
    atTopLevel (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
unop_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 (String -> Name
nameFromString (String -> Name) -> (PrimType -> String) -> PrimType -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> String
forall a. Pretty a => a -> String
pretty) [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
              (String -> Name
nameFromString (String -> Name) -> (BinOp -> String) -> BinOp -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> String
forall a. Pretty a => a -> String
pretty)
              [BinOp
forall a. Bounded a => a
minBound .. (BinOp
forall a. Bounded a => a
maxBound :: BinOp)]
        unop_names :: Set Name
unop_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
$ (String -> Name) -> [String] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map String -> Name
nameFromString [String
"!"]
        fun_names :: Set Name
fun_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
$ (String -> Name) -> [String] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map String -> Name
nameFromString [String
"shape"]
    atTopLevel (Namespace, Name)
_ = Bool
False