{-# LANGUAGE Strict #-}

-- | Facilities for type-checking terms.  Factored out of
-- "Language.Futhark.TypeChecker.Terms" to prevent the module from
-- being gigantic.
--
-- Incidentally also a nice place to put Haddock comments to make the
-- internal API of the type checker easier to browse.
module Language.Futhark.TypeChecker.Terms.Monad
  ( TermTypeM,
    runTermTypeM,
    ValBinding (..),
    SizeSource (SourceBound, SourceSlice),
    Inferred (..),
    Checking (..),
    withEnv,
    localScope,
    TermEnv (..),
    TermScope (..),
    TermTypeState (..),
    onFailure,
    extSize,
    expType,
    expTypeFully,
    constrain,
    newArrayType,
    allDimsFreshInType,
    updateTypes,
    Names,

    -- * Primitive checking
    unifies,
    require,
    checkTypeExpNonrigid,

    -- * Sizes
    isInt64,

    -- * Control flow
    incLevel,

    -- * Errors
    unusedSize,
  )
where

import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bitraversable
import Data.Char (isAscii)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.FreshNames hiding (newName)
import Futhark.FreshNames qualified
import Futhark.Util.Pretty hiding (space)
import Language.Futhark
import Language.Futhark.Semantic (includeToFilePath)
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Monad hiding (BoundV, stateNameSource)
import Language.Futhark.TypeChecker.Monad qualified as TypeM
import Language.Futhark.TypeChecker.Types
import Language.Futhark.TypeChecker.Unify
import Prelude hiding (mod)

type Names = S.Set VName

data ValBinding
  = BoundV [TypeParam] StructType
  | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType)
  | EqualityF
  deriving (Int -> ValBinding -> ShowS
[ValBinding] -> ShowS
ValBinding -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ValBinding] -> ShowS
$cshowList :: [ValBinding] -> ShowS
show :: ValBinding -> String
$cshow :: ValBinding -> String
showsPrec :: Int -> ValBinding -> ShowS
$cshowsPrec :: Int -> ValBinding -> ShowS
Show)

unusedSize :: (MonadTypeChecker m) => SizeBinder VName -> m a
unusedSize :: forall (m :: * -> *) a.
MonadTypeChecker m =>
SizeBinder VName -> m a
unusedSize SizeBinder VName
p =
  forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SizeBinder VName
p forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"unused-size" forall a b. (a -> b) -> a -> b
$
    Doc ()
"Size" forall a. Doc a -> Doc a -> Doc a
<+> forall a ann. Pretty a => a -> Doc ann
pretty SizeBinder VName
p forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"unused in pattern."

data Inferred t
  = NoneInferred
  | Ascribed t

instance Functor Inferred where
  fmap :: forall a b. (a -> b) -> Inferred a -> Inferred b
fmap a -> b
_ Inferred a
NoneInferred = forall t. Inferred t
NoneInferred
  fmap a -> b
f (Ascribed a
t) = forall t. t -> Inferred t
Ascribed (a -> b
f a
t)

data Checking
  = CheckingApply (Maybe (QualName VName)) Exp StructType StructType
  | CheckingReturn ResType StructType
  | CheckingAscription StructType StructType
  | CheckingLetGeneralise Name
  | CheckingParams (Maybe Name)
  | CheckingPat (UncheckedPat StructType) (Inferred StructType)
  | CheckingLoopBody StructType StructType
  | CheckingLoopInitial StructType StructType
  | CheckingRecordUpdate [Name] StructType StructType
  | CheckingRequired [StructType] StructType
  | CheckingBranches StructType StructType

instance Pretty Checking where
  pretty :: forall ann. Checking -> Doc ann
pretty (CheckingApply Maybe (QualName VName)
f Exp
e StructType
expected StructType
actual) =
    forall {ann}. Doc ann
header
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Expected:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
expected)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual:  "
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
    where
      header :: Doc ann
header =
        case Maybe (QualName VName)
f of
          Maybe (QualName VName)
Nothing ->
            Doc ann
"Cannot apply function to"
              forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a b. Doc a -> Doc b
shorten forall a b. (a -> b) -> a -> b
$ forall ann. Doc ann -> Doc ann
group forall a b. (a -> b) -> a -> b
$ forall a ann. Pretty a => a -> Doc ann
pretty Exp
e) forall a. Semigroup a => a -> a -> a
<> Doc ann
" (invalid type)."
          Just QualName VName
fname ->
            Doc ann
"Cannot apply"
              forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName VName
fname)
              forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
"to"
              forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes (forall ann. Doc ann -> Doc ann
align forall a b. (a -> b) -> a -> b
$ forall a b. Doc a -> Doc b
shorten forall a b. (a -> b) -> a -> b
$ forall ann. Doc ann -> Doc ann
group forall a b. (a -> b) -> a -> b
$ forall a ann. Pretty a => a -> Doc ann
pretty Exp
e) forall a. Semigroup a => a -> a -> a
<> Doc ann
" (invalid type)."
  pretty (CheckingReturn ResType
expected StructType
actual) =
    Doc ann
"Function body does not have expected type."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Expected:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty ResType
expected)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual:  "
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
  pretty (CheckingAscription StructType
expected StructType
actual) =
    Doc ann
"Expression does not have expected type from explicit ascription."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Expected:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
expected)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual:  "
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
  pretty (CheckingLetGeneralise Name
fname) =
    Doc ann
"Cannot generalise type of" forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty Name
fname) forall a. Semigroup a => a -> a -> a
<> Doc ann
"."
  pretty (CheckingParams Maybe Name
fname) =
    Doc ann
"Invalid use of parameters in" forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes forall {ann}. Doc ann
fname' forall a. Semigroup a => a -> a -> a
<> Doc ann
"."
    where
      fname' :: Doc ann
fname' = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Doc ann
"anonymous function" forall a ann. Pretty a => a -> Doc ann
pretty Maybe Name
fname
  pretty (CheckingPat UncheckedPat StructType
pat Inferred StructType
NoneInferred) =
    Doc ann
"Invalid pattern" forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty UncheckedPat StructType
pat) forall a. Semigroup a => a -> a -> a
<> Doc ann
"."
  pretty (CheckingPat UncheckedPat StructType
pat (Ascribed StructType
t)) =
    Doc ann
"Pattern"
      forall a. Doc a -> Doc a -> Doc a
</> forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall a ann. Pretty a => a -> Doc ann
pretty UncheckedPat StructType
pat)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"cannot match value of type"
      forall a. Doc a -> Doc a -> Doc a
</> forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall a ann. Pretty a => a -> Doc ann
pretty StructType
t)
  pretty (CheckingLoopBody StructType
expected StructType
actual) =
    Doc ann
"Loop body does not have expected type."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Expected:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
expected)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual:  "
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
  pretty (CheckingLoopInitial StructType
expected StructType
actual) =
    Doc ann
"Initial loop values do not have expected type."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Expected:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
expected)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual:  "
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
  pretty (CheckingRecordUpdate [Name]
fs StructType
expected StructType
actual) =
    Doc ann
"Type mismatch when updating record field"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes forall {ann}. Doc ann
fs' forall a. Semigroup a => a -> a -> a
<> Doc ann
"."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Existing:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
expected)
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"New:     "
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
    where
      fs' :: Doc ann
fs' = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall ann. Doc ann -> [Doc ann] -> [Doc ann]
punctuate Doc ann
"." forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a ann. Pretty a => a -> Doc ann
pretty [Name]
fs
  pretty (CheckingRequired [StructType
expected] StructType
actual) =
    Doc ann
"Expression must must have type"
      forall a. Doc a -> Doc a -> Doc a
<+> forall a ann. Pretty a => a -> Doc ann
pretty StructType
expected forall a. Semigroup a => a -> a -> a
<> Doc ann
"."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual type:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
  pretty (CheckingRequired [StructType]
expected StructType
actual) =
    Doc ann
"Type of expression must must be one of "
      forall a. Doc a -> Doc a -> Doc a
<+> forall {ann}. Doc ann
expected' forall a. Semigroup a => a -> a -> a
<> Doc ann
"."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Actual type:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty StructType
actual)
    where
      expected' :: Doc a
expected' = forall a. [Doc a] -> Doc a
commasep (forall a b. (a -> b) -> [a] -> [b]
map forall a ann. Pretty a => a -> Doc ann
pretty [StructType]
expected)
  pretty (CheckingBranches StructType
t1 StructType
t2) =
    Doc ann
"Branches differ in type."
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Former:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall a ann. Pretty a => a -> Doc ann
pretty StructType
t1
      forall a. Doc a -> Doc a -> Doc a
</> Doc ann
"Latter:"
      forall a. Doc a -> Doc a -> Doc a
<+> forall a ann. Pretty a => a -> Doc ann
pretty StructType
t2

-- | Type checking happens with access to this environment.  The
-- 'TermScope' will be extended during type-checking as bindings come into
-- scope.
data TermEnv = TermEnv
  { TermEnv -> TermScope
termScope :: TermScope,
    TermEnv -> Maybe Checking
termChecking :: Maybe Checking,
    TermEnv -> Int
termLevel :: Level,
    TermEnv -> UncheckedExp -> TermTypeM Exp
termChecker :: UncheckedExp -> TermTypeM Exp,
    TermEnv -> Env
termOuterEnv :: Env,
    TermEnv -> ImportName
termImportName :: ImportName
  }

data TermScope = TermScope
  { TermScope -> Map VName ValBinding
scopeVtable :: M.Map VName ValBinding,
    TermScope -> Map VName TypeBinding
scopeTypeTable :: M.Map VName TypeBinding,
    TermScope -> Map VName Mod
scopeModTable :: M.Map VName Mod,
    TermScope -> NameMap
scopeNameMap :: NameMap
  }
  deriving (Int -> TermScope -> ShowS
[TermScope] -> ShowS
TermScope -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TermScope] -> ShowS
$cshowList :: [TermScope] -> ShowS
show :: TermScope -> String
$cshow :: TermScope -> String
showsPrec :: Int -> TermScope -> ShowS
$cshowsPrec :: Int -> TermScope -> ShowS
Show)

instance Semigroup TermScope where
  TermScope Map VName ValBinding
vt1 Map VName TypeBinding
tt1 Map VName Mod
mt1 NameMap
nt1 <> :: TermScope -> TermScope -> TermScope
<> TermScope Map VName ValBinding
vt2 Map VName TypeBinding
tt2 Map VName Mod
mt2 NameMap
nt2 =
    Map VName ValBinding
-> Map VName TypeBinding -> Map VName Mod -> NameMap -> TermScope
TermScope (Map VName ValBinding
vt2 forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName ValBinding
vt1) (Map VName TypeBinding
tt2 forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName TypeBinding
tt1) (Map VName Mod
mt1 forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName Mod
mt2) (NameMap
nt2 forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` NameMap
nt1)

envToTermScope :: Env -> TermScope
envToTermScope :: Env -> TermScope
envToTermScope Env
env =
  TermScope
    { scopeVtable :: Map VName ValBinding
scopeVtable = Map VName ValBinding
vtable,
      scopeTypeTable :: Map VName TypeBinding
scopeTypeTable = Env -> Map VName TypeBinding
envTypeTable Env
env,
      scopeNameMap :: NameMap
scopeNameMap = Env -> NameMap
envNameMap Env
env,
      scopeModTable :: Map VName Mod
scopeModTable = Env -> Map VName Mod
envModTable Env
env
    }
  where
    vtable :: Map VName ValBinding
vtable = forall a b k. (a -> b) -> Map k a -> Map k b
M.map BoundV -> ValBinding
valBinding forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable Env
env
    valBinding :: BoundV -> ValBinding
valBinding (TypeM.BoundV [TypeParam]
tps StructType
v) = [TypeParam] -> StructType -> ValBinding
BoundV [TypeParam]
tps StructType
v

withEnv :: TermEnv -> Env -> TermEnv
withEnv :: TermEnv -> Env -> TermEnv
withEnv TermEnv
tenv Env
env = TermEnv
tenv {termScope :: TermScope
termScope = TermEnv -> TermScope
termScope TermEnv
tenv forall a. Semigroup a => a -> a -> a
<> Env -> TermScope
envToTermScope Env
env}

-- | Wrap a function name to give it a vacuous Eq instance for SizeSource.
newtype FName = FName (Maybe (QualName VName))
  deriving (Int -> FName -> ShowS
[FName] -> ShowS
FName -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FName] -> ShowS
$cshowList :: [FName] -> ShowS
show :: FName -> String
$cshow :: FName -> String
showsPrec :: Int -> FName -> ShowS
$cshowsPrec :: Int -> FName -> ShowS
Show)

instance Eq FName where
  FName
_ == :: FName -> FName -> Bool
== FName
_ = Bool
True

instance Ord FName where
  compare :: FName -> FName -> Ordering
compare FName
_ FName
_ = Ordering
EQ

-- | What was the source of some existential size?  This is used for
-- using the same existential variable if the same source is
-- encountered in multiple locations.
data SizeSource
  = SourceArg FName (ExpBase NoInfo VName)
  | SourceBound (ExpBase NoInfo VName)
  | SourceSlice
      (Maybe Size)
      (Maybe (ExpBase NoInfo VName))
      (Maybe (ExpBase NoInfo VName))
      (Maybe (ExpBase NoInfo VName))
  deriving (SizeSource -> SizeSource -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SizeSource -> SizeSource -> Bool
$c/= :: SizeSource -> SizeSource -> Bool
== :: SizeSource -> SizeSource -> Bool
$c== :: SizeSource -> SizeSource -> Bool
Eq, Eq SizeSource
SizeSource -> SizeSource -> Bool
SizeSource -> SizeSource -> Ordering
SizeSource -> SizeSource -> SizeSource
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: SizeSource -> SizeSource -> SizeSource
$cmin :: SizeSource -> SizeSource -> SizeSource
max :: SizeSource -> SizeSource -> SizeSource
$cmax :: SizeSource -> SizeSource -> SizeSource
>= :: SizeSource -> SizeSource -> Bool
$c>= :: SizeSource -> SizeSource -> Bool
> :: SizeSource -> SizeSource -> Bool
$c> :: SizeSource -> SizeSource -> Bool
<= :: SizeSource -> SizeSource -> Bool
$c<= :: SizeSource -> SizeSource -> Bool
< :: SizeSource -> SizeSource -> Bool
$c< :: SizeSource -> SizeSource -> Bool
compare :: SizeSource -> SizeSource -> Ordering
$ccompare :: SizeSource -> SizeSource -> Ordering
Ord, Int -> SizeSource -> ShowS
[SizeSource] -> ShowS
SizeSource -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SizeSource] -> ShowS
$cshowList :: [SizeSource] -> ShowS
show :: SizeSource -> String
$cshow :: SizeSource -> String
showsPrec :: Int -> SizeSource -> ShowS
$cshowsPrec :: Int -> SizeSource -> ShowS
Show)

-- | The state is a set of constraints and a counter for generating
-- type names.  This is distinct from the usual counter we use for
-- generating unique names, as these will be user-visible.
data TermTypeState = TermTypeState
  { TermTypeState -> Constraints
stateConstraints :: Constraints,
    TermTypeState -> Int
stateCounter :: !Int,
    TermTypeState -> Set VName
stateUsed :: S.Set VName,
    TermTypeState -> Warnings
stateWarnings :: Warnings,
    TermTypeState -> VNameSource
stateNameSource :: VNameSource
  }

newtype TermTypeM a
  = TermTypeM
      ( ReaderT
          TermEnv
          (StateT TermTypeState (Except (Warnings, TypeError)))
          a
      )
  deriving
    ( Applicative TermTypeM
forall a. a -> TermTypeM a
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall a b. TermTypeM a -> (a -> TermTypeM b) -> TermTypeM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> TermTypeM a
$creturn :: forall a. a -> TermTypeM a
>> :: forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
$c>> :: forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
>>= :: forall a b. TermTypeM a -> (a -> TermTypeM b) -> TermTypeM b
$c>>= :: forall a b. TermTypeM a -> (a -> TermTypeM b) -> TermTypeM b
Monad,
      forall a b. a -> TermTypeM b -> TermTypeM a
forall a b. (a -> b) -> TermTypeM a -> TermTypeM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> TermTypeM b -> TermTypeM a
$c<$ :: forall a b. a -> TermTypeM b -> TermTypeM a
fmap :: forall a b. (a -> b) -> TermTypeM a -> TermTypeM b
$cfmap :: forall a b. (a -> b) -> TermTypeM a -> TermTypeM b
Functor,
      Functor TermTypeM
forall a. a -> TermTypeM a
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM a
forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
forall a b c.
(a -> b -> c) -> TermTypeM a -> TermTypeM b -> TermTypeM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. TermTypeM a -> TermTypeM b -> TermTypeM a
$c<* :: forall a b. TermTypeM a -> TermTypeM b -> TermTypeM a
*> :: forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
$c*> :: forall a b. TermTypeM a -> TermTypeM b -> TermTypeM b
liftA2 :: forall a b c.
(a -> b -> c) -> TermTypeM a -> TermTypeM b -> TermTypeM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> TermTypeM a -> TermTypeM b -> TermTypeM c
<*> :: forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
$c<*> :: forall a b. TermTypeM (a -> b) -> TermTypeM a -> TermTypeM b
pure :: forall a. a -> TermTypeM a
$cpure :: forall a. a -> TermTypeM a
Applicative,
      MonadReader TermEnv,
      MonadState TermTypeState
    )

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

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

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

constrain :: VName -> Constraint -> TermTypeM ()
constrain :: VName -> Constraint -> TermTypeM ()
constrain VName
v Constraint
c = do
  Int
lvl <- forall (m :: * -> *). MonadUnify m => m Int
curLevel
  forall (m :: * -> *).
MonadUnify m =>
(Constraints -> Constraints) -> m ()
modifyConstraints forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Int
lvl, Constraint
c)

instance MonadUnify TermTypeM where
  getConstraints :: TermTypeM Constraints
getConstraints = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TermTypeState -> Constraints
stateConstraints
  putConstraints :: Constraints -> TermTypeM ()
putConstraints Constraints
x = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateConstraints :: Constraints
stateConstraints = Constraints
x}

  newTypeVar :: forall als dim.
Monoid als =>
SrcLoc -> Name -> TermTypeM (TypeBase dim als)
newTypeVar SrcLoc
loc Name
desc = do
    Int
i <- TermTypeM Int
incCounter
    VName
v <- forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID forall a b. (a -> b) -> a -> b
$ Name -> Int -> Name
mkTypeVarName Name
desc Int
i
    VName -> Constraint -> TermTypeM ()
constrain VName
v forall a b. (a -> b) -> a -> b
$ Liftedness -> Usage -> Constraint
NoConstraint Liftedness
Lifted forall a b. (a -> b) -> a -> b
$ forall a. Located a => a -> Usage
mkUsage' SrcLoc
loc
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar forall a. Monoid a => a
mempty (forall v. v -> QualName v
qualName VName
v) []

  curLevel :: TermTypeM Int
curLevel = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Int
termLevel

  newDimVar :: Usage -> Rigidity -> Name -> TermTypeM VName
newDimVar Usage
usage Rigidity
rigidity Name
name = do
    VName
dim <- forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
name
    case Rigidity
rigidity of
      Rigid RigidSource
rsrc -> VName -> Constraint -> TermTypeM ()
constrain VName
dim forall a b. (a -> b) -> a -> b
$ SrcLoc -> RigidSource -> Constraint
UnknownSize (forall a. Located a => a -> SrcLoc
srclocOf Usage
usage) RigidSource
rsrc
      Rigidity
Nonrigid -> VName -> Constraint -> TermTypeM ()
constrain VName
dim forall a b. (a -> b) -> a -> b
$ Maybe Exp -> Usage -> Constraint
Size forall a. Maybe a
Nothing Usage
usage
    forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
dim

  unifyError :: forall loc a.
Located loc =>
loc -> Notes -> BreadCrumbs -> Doc () -> TermTypeM a
unifyError loc
loc Notes
notes BreadCrumbs
bcs Doc ()
doc = do
    Maybe Checking
checking <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Maybe Checking
termChecking
    case Maybe Checking
checking of
      Just Checking
checking' ->
        forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$
          Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes forall a b. (a -> b) -> a -> b
$
            forall a ann. Pretty a => a -> Doc ann
pretty Checking
checking' forall a. Semigroup a => a -> a -> a
<> forall {ann}. Doc ann
line forall a. Doc a -> Doc a -> Doc a
</> Doc ()
doc forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty BreadCrumbs
bcs
      Maybe Checking
Nothing ->
        forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$ Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes forall a b. (a -> b) -> a -> b
$ Doc ()
doc forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty BreadCrumbs
bcs

  matchError :: forall loc a.
Located loc =>
loc
-> Notes -> BreadCrumbs -> StructType -> StructType -> TermTypeM a
matchError loc
loc Notes
notes BreadCrumbs
bcs StructType
t1 StructType
t2 = do
    Maybe Checking
checking <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Maybe Checking
termChecking
    case Maybe Checking
checking of
      Just Checking
checking'
        | BreadCrumbs -> Bool
hasNoBreadCrumbs BreadCrumbs
bcs ->
            forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$
              Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes forall a b. (a -> b) -> a -> b
$
                forall a ann. Pretty a => a -> Doc ann
pretty Checking
checking'
        | Bool
otherwise ->
            forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$
              Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes forall a b. (a -> b) -> a -> b
$
                forall a ann. Pretty a => a -> Doc ann
pretty Checking
checking' forall a. Semigroup a => a -> a -> a
<> forall {ann}. Doc ann
line forall a. Doc a -> Doc a -> Doc a
</> forall {ann}. Doc ann
doc forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty BreadCrumbs
bcs
      Maybe Checking
Nothing ->
        forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$ Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes forall a b. (a -> b) -> a -> b
$ forall {ann}. Doc ann
doc forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty BreadCrumbs
bcs
    where
      doc :: Doc a
doc =
        Doc a
"Types"
          forall a. Doc a -> Doc a -> Doc a
</> forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall a ann. Pretty a => a -> Doc ann
pretty StructType
t1)
          forall a. Doc a -> Doc a -> Doc a
</> Doc a
"and"
          forall a. Doc a -> Doc a -> Doc a
</> forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall a ann. Pretty a => a -> Doc ann
pretty StructType
t2)
          forall a. Doc a -> Doc a -> Doc a
</> Doc a
"do not match."

-- | Instantiate a type scheme with fresh type variables for its type
-- parameters. Returns the names of the fresh type variables, the
-- instance list, and the instantiated type.
instantiateTypeScheme ::
  QualName VName ->
  SrcLoc ->
  [TypeParam] ->
  StructType ->
  TermTypeM ([VName], StructType)
instantiateTypeScheme :: QualName VName
-> SrcLoc
-> [TypeParam]
-> StructType
-> TermTypeM ([VName], StructType)
instantiateTypeScheme QualName VName
qn SrcLoc
loc [TypeParam]
tparams StructType
t = do
  let tnames :: [VName]
tnames = forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
tparams
  ([VName]
tparam_names, [Subst (RetTypeBase Exp NoUniqueness)]
tparam_substs) <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM (forall as dim.
Monoid as =>
QualName VName
-> SrcLoc
-> TypeParam
-> TermTypeM (VName, Subst (RetTypeBase dim as))
instantiateTypeParam QualName VName
qn SrcLoc
loc) [TypeParam]
tparams
  let substs :: Map VName (Subst (RetTypeBase Exp NoUniqueness))
substs = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
tnames [Subst (RetTypeBase Exp NoUniqueness)]
tparam_substs
      t' :: StructType
t' = forall a. Substitutable a => TypeSubs -> a -> a
applySubst (forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (RetTypeBase Exp NoUniqueness))
substs) StructType
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName]
tparam_names, StructType
t')

-- | Create a new type name and insert it (unconstrained) in the
-- substitution map.
instantiateTypeParam ::
  Monoid as =>
  QualName VName ->
  SrcLoc ->
  TypeParam ->
  TermTypeM (VName, Subst (RetTypeBase dim as))
instantiateTypeParam :: forall as dim.
Monoid as =>
QualName VName
-> SrcLoc
-> TypeParam
-> TermTypeM (VName, Subst (RetTypeBase dim as))
instantiateTypeParam QualName VName
qn SrcLoc
loc TypeParam
tparam = do
  Int
i <- TermTypeM Int
incCounter
  let name :: Name
name = String -> Name
nameFromString (forall a. (a -> Bool) -> [a] -> [a]
takeWhile Char -> Bool
isAscii (VName -> String
baseString (forall vn. TypeParamBase vn -> vn
typeParamName TypeParam
tparam)))
  VName
v <- forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID forall a b. (a -> b) -> a -> b
$ Name -> Int -> Name
mkTypeVarName Name
name Int
i
  case TypeParam
tparam of
    TypeParamType Liftedness
x VName
_ SrcLoc
_ -> do
      VName -> Constraint -> TermTypeM ()
constrain VName
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. Liftedness -> Usage -> Constraint
NoConstraint Liftedness
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
        Doc Any
"instantiated type parameter of " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName VName
qn)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, forall t. [TypeParam] -> t -> Subst t
Subst [] forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar forall a. Monoid a => a
mempty (forall v. v -> QualName v
qualName VName
v) [])
    TypeParamDim {} -> do
      VName -> Constraint -> TermTypeM ()
constrain VName
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Exp -> Usage -> Constraint
Size forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
        Doc Any
"instantiated size parameter of " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName VName
qn)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, forall t. Exp -> Subst t
ExpSubst forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (forall v. v -> QualName v
qualName VName
v) SrcLoc
loc)

checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkQualNameWithEnv :: Namespace
-> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkQualNameWithEnv Namespace
space qn :: QualName Name
qn@(QualName [Name]
quals Name
name) SrcLoc
loc = do
  TermScope
scope <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> TermScope
termScope
  TermScope -> [Name] -> TermTypeM (TermScope, QualName VName)
descend TermScope
scope [Name]
quals
  where
    descend :: TermScope -> [Name] -> TermTypeM (TermScope, QualName VName)
descend TermScope
scope []
      | Just QualName VName
name' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
space, Name
name) forall a b. (a -> b) -> a -> b
$ TermScope -> NameMap
scopeNameMap TermScope
scope =
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (TermScope
scope, QualName VName
name')
      | Bool
otherwise =
          forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
qn SrcLoc
loc
    descend TermScope
scope (Name
q : [Name]
qs)
      | Just (QualName [VName]
_ VName
q') <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
Term, Name
q) forall a b. (a -> b) -> a -> b
$ TermScope -> NameMap
scopeNameMap TermScope
scope,
        Just Mod
res <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q' forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName Mod
scopeModTable TermScope
scope =
          case Mod
res of
            -- Check if we are referring to the magical intrinsics
            -- module.
            Mod
_
              | VName -> Int
baseTag VName
q' forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag ->
                  Namespace
-> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkIntrinsic Namespace
space QualName Name
qn SrcLoc
loc
            ModEnv Env
q_scope -> do
              (TermScope
scope', QualName [VName]
qs' VName
name') <- TermScope -> [Name] -> TermTypeM (TermScope, QualName VName)
descend (Env -> TermScope
envToTermScope Env
q_scope) [Name]
qs
              forall (f :: * -> *) a. Applicative f => a -> f a
pure (TermScope
scope', forall vn. [vn] -> vn -> QualName vn
QualName (VName
q' forall a. a -> [a] -> [a]
: [VName]
qs') VName
name')
            ModFun {} -> forall (m :: * -> *) a. MonadTypeChecker m => SrcLoc -> m a
unappliedFunctor SrcLoc
loc
      | Bool
otherwise =
          forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
qn SrcLoc
loc

checkIntrinsic :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkIntrinsic :: Namespace
-> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkIntrinsic Namespace
space qn :: QualName Name
qn@(QualName [Name]
_ Name
name) SrcLoc
loc
  | Just QualName VName
v <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
space, Name
name) NameMap
intrinsicsNameMap = do
      ImportName
me <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> ImportName
termImportName
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (String -> Bool
isBuiltin (ImportName -> String
includeToFilePath ImportName
me)) forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn SrcLoc
loc Doc ()
"Using intrinsic functions directly can easily crash the compiler or result in wrong code generation."
      TermScope
scope <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> TermScope
termScope
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (TermScope
scope, QualName VName
v)
  | Bool
otherwise =
      forall (m :: * -> *) a.
MonadTypeChecker m =>
Namespace -> QualName Name -> SrcLoc -> m a
unknownVariable Namespace
space QualName Name
qn SrcLoc
loc

localScope :: (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope :: forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope TermScope -> TermScope
f = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \TermEnv
tenv -> TermEnv
tenv {termScope :: TermScope
termScope = TermScope -> TermScope
f forall a b. (a -> b) -> a -> b
$ TermEnv -> TermScope
termScope TermEnv
tenv}

instance MonadTypeChecker TermTypeM where
  checkExpForSize :: UncheckedExp -> TermTypeM Exp
checkExpForSize UncheckedExp
e = do
    UncheckedExp -> TermTypeM Exp
checker <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> UncheckedExp -> TermTypeM Exp
termChecker
    Exp
e' <- UncheckedExp -> TermTypeM Exp
checker UncheckedExp
e
    let t :: StructType
t = forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e'
    forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (forall a. Located a => a -> Text -> Usage
mkUsage (forall a. Located a => a -> SrcLoc
srclocOf Exp
e') Text
"Size expression") StructType
t (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (forall dim u. PrimType -> ScalarTypeBase dim u
Prim (IntType -> PrimType
Signed IntType
Int64)))
    forall e. ASTMappable e => e -> TermTypeM e
updateTypes Exp
e'

  warnings :: Warnings -> TermTypeM ()
warnings Warnings
ws =
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateWarnings :: Warnings
stateWarnings = TermTypeState -> Warnings
stateWarnings TermTypeState
s forall a. Semigroup a => a -> a -> a
<> Warnings
ws}

  warn :: forall loc. Located loc => loc -> Doc () -> TermTypeM ()
warn loc
loc Doc ()
problem = forall (m :: * -> *). MonadTypeChecker m => Warnings -> m ()
warnings forall a b. (a -> b) -> a -> b
$ SrcLoc -> Doc () -> Warnings
singleWarning (forall a. Located a => a -> SrcLoc
srclocOf loc
loc) Doc ()
problem

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

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

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

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

  bindNameMap :: forall a. NameMap -> TermTypeM a -> TermTypeM a
bindNameMap NameMap
m = forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope forall a b. (a -> b) -> a -> b
$ \TermScope
scope ->
    TermScope
scope {scopeNameMap :: NameMap
scopeNameMap = NameMap
m forall a. Semigroup a => a -> a -> a
<> TermScope -> NameMap
scopeNameMap TermScope
scope}

  bindVal :: forall a. VName -> BoundV -> TermTypeM a -> TermTypeM a
bindVal VName
v (TypeM.BoundV [TypeParam]
tps StructType
t) = forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope forall a b. (a -> b) -> a -> b
$ \TermScope
scope ->
    TermScope
scope {scopeVtable :: Map VName ValBinding
scopeVtable = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v ([TypeParam] -> StructType -> ValBinding
BoundV [TypeParam]
tps StructType
t) forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope}

  lookupType :: SrcLoc
-> QualName Name
-> TermTypeM
     (QualName VName, [TypeParam], RetTypeBase Exp NoUniqueness,
      Liftedness)
lookupType SrcLoc
loc QualName Name
qn = do
    Env
outer_env <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Env
termOuterEnv
    (TermScope
scope, qn' :: QualName VName
qn'@(QualName [VName]
qs VName
name)) <- Namespace
-> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkQualNameWithEnv Namespace
Type QualName Name
qn SrcLoc
loc
    case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName TypeBinding
scopeTypeTable TermScope
scope of
      Maybe TypeBinding
Nothing -> forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
unknownType SrcLoc
loc QualName Name
qn
      Just (TypeAbbr Liftedness
l [TypeParam]
ps (RetType [VName]
dims StructType
def)) ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure
          ( QualName VName
qn',
            [TypeParam]
ps,
            forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims forall a b. (a -> b) -> a -> b
$ forall as.
Env -> [VName] -> [VName] -> TypeBase Exp as -> TypeBase Exp as
qualifyTypeVars Env
outer_env (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
ps) [VName]
qs StructType
def,
            Liftedness
l
          )

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

  lookupVar :: SrcLoc -> QualName Name -> TermTypeM (QualName VName, StructType)
lookupVar SrcLoc
loc QualName Name
qn = do
    (TermScope
scope, qn' :: QualName VName
qn'@(QualName [VName]
qs VName
name)) <- Namespace
-> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkQualNameWithEnv Namespace
Term QualName Name
qn SrcLoc
loc
    let usage :: Usage
usage = forall a. Located a => a -> Text -> Usage
mkUsage SrcLoc
loc forall a b. (a -> b) -> a -> b
$ forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$ Doc Any
"use of " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName Name
qn)

    StructType
t <- case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope of
      Maybe ValBinding
Nothing ->
        forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError SrcLoc
loc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
          Doc ()
"Unknown variable" forall a. Doc a -> Doc a -> Doc a
<+> forall ann. Doc ann -> Doc ann
dquotes (forall a ann. Pretty a => a -> Doc ann
pretty QualName Name
qn) forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
      Just (BoundV [TypeParam]
tparams StructType
t) -> do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
qs) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \TermTypeState
s ->
          TermTypeState
s {stateUsed :: Set VName
stateUsed = forall a. Ord a => a -> Set a -> Set a
S.insert (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn') forall a b. (a -> b) -> a -> b
$ TermTypeState -> Set VName
stateUsed TermTypeState
s}
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Text -> Char
T.head (Name -> Text
nameToText (VName -> Name
baseName VName
name)) forall a. Eq a => a -> a -> Bool
== Char
'_') forall a b. (a -> b) -> a -> b
$
          forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
underscoreUse SrcLoc
loc QualName Name
qn
        if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParam]
tparams Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
qs
          then forall (f :: * -> *) a. Applicative f => a -> f a
pure StructType
t
          else do
            ([VName]
tnames, StructType
t') <- QualName VName
-> SrcLoc
-> [TypeParam]
-> StructType
-> TermTypeM ([VName], StructType)
instantiateTypeScheme QualName VName
qn' SrcLoc
loc [TypeParam]
tparams StructType
t
            Env
outer_env <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Env
termOuterEnv
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall as.
Env -> [VName] -> [VName] -> TypeBase Exp as -> TypeBase Exp as
qualifyTypeVars Env
outer_env [VName]
tnames [VName]
qs StructType
t'
      Just ValBinding
EqualityF -> do
        StructType
argtype <- forall (m :: * -> *) als dim.
(MonadUnify m, Monoid als) =>
SrcLoc -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
        forall (m :: * -> *) dim u.
(MonadUnify m, Pretty (Shape dim), Pretty u) =>
Usage -> TypeBase dim u -> m ()
equalityType Usage
usage StructType
argtype
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
          forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe StructType
argtype forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$
            forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$
              forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe StructType
argtype forall a b. (a -> b) -> a -> b
$
                forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$
                  forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$
                    forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
Bool
      Just (OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt) -> do
        StructType
argtype <- forall (m :: * -> *) als dim.
(MonadUnify m, Monoid als) =>
SrcLoc -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
        forall (m :: * -> *).
MonadUnify m =>
[PrimType] -> Usage -> StructType -> m ()
mustBeOneOf [PrimType]
ts Usage
usage StructType
argtype
        let ([StructType]
pts', StructType
rt') = forall {dim} {u}.
TypeBase dim u
-> [Maybe PrimType]
-> Maybe PrimType
-> ([TypeBase dim NoUniqueness], TypeBase dim NoUniqueness)
instOverloaded StructType
argtype [Maybe PrimType]
pts Maybe PrimType
rt
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ParamType] -> RetTypeBase Exp Uniqueness -> StructType
foldFunType (forall a b. (a -> b) -> [a] -> [b]
map (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe) [StructType]
pts') forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Exp u -> ResType
toRes Uniqueness
Nonunique StructType
rt'

    forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName
qn', StructType
t)
    where
      instOverloaded :: TypeBase dim u
-> [Maybe PrimType]
-> Maybe PrimType
-> ([TypeBase dim NoUniqueness], TypeBase dim NoUniqueness)
instOverloaded TypeBase dim u
argtype [Maybe PrimType]
pts Maybe PrimType
rt =
        ( forall a b. (a -> b) -> [a] -> [b]
map (forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase dim u
argtype) (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u. PrimType -> ScalarTypeBase dim u
Prim)) [Maybe PrimType]
pts,
          forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase dim u
argtype) (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u. PrimType -> ScalarTypeBase dim u
Prim) Maybe PrimType
rt
        )

  typeError :: forall loc a. Located loc => loc -> Notes -> Doc () -> TermTypeM a
typeError loc
loc Notes
notes Doc ()
s = do
    Maybe Checking
checking <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Maybe Checking
termChecking
    case Maybe Checking
checking of
      Just Checking
checking' ->
        forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$ Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes (forall a ann. Pretty a => a -> Doc ann
pretty Checking
checking' forall a. Semigroup a => a -> a -> a
<> forall {ann}. Doc ann
line forall a. Doc a -> Doc a -> Doc a
</> Doc ()
s)
      Maybe Checking
Nothing ->
        forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall a b. (a -> b) -> a -> b
$ Loc -> Notes -> Doc () -> TypeError
TypeError (forall a. Located a => a -> Loc
locOf loc
loc) Notes
notes Doc ()
s

onFailure :: Checking -> TermTypeM a -> TermTypeM a
onFailure :: forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure Checking
c = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \TermEnv
env -> TermEnv
env {termChecking :: Maybe Checking
termChecking = forall a. a -> Maybe a
Just Checking
c}

extSize :: SrcLoc -> SizeSource -> TermTypeM (Size, Maybe VName)
extSize :: SrcLoc -> SizeSource -> TermTypeM (Exp, Maybe VName)
extSize SrcLoc
loc SizeSource
e = do
  let rsrc :: RigidSource
rsrc = case SizeSource
e of
        SourceArg (FName Maybe (QualName VName)
fname) ExpBase NoInfo VName
e' ->
          Maybe (QualName VName) -> Text -> RigidSource
RigidArg Maybe (QualName VName)
fname forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyTextOneLine ExpBase NoInfo VName
e'
        SourceBound ExpBase NoInfo VName
e' ->
          Text -> RigidSource
RigidBound forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyTextOneLine ExpBase NoInfo VName
e'
        SourceSlice Maybe Exp
d Maybe (ExpBase NoInfo VName)
i Maybe (ExpBase NoInfo VName)
j Maybe (ExpBase NoInfo VName)
s ->
          Maybe Exp -> Text -> RigidSource
RigidSlice Maybe Exp
d forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyTextOneLine forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice Maybe (ExpBase NoInfo VName)
i Maybe (ExpBase NoInfo VName)
j Maybe (ExpBase NoInfo VName)
s
  VName
d <- forall (m :: * -> *) a.
(MonadUnify m, Located a) =>
a -> RigidSource -> Name -> m VName
newRigidDim SrcLoc
loc RigidSource
rsrc Name
"n"
  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( QualName VName -> SrcLoc -> Exp
sizeFromName (forall v. v -> QualName v
qualName VName
d) SrcLoc
loc,
      forall a. a -> Maybe a
Just VName
d
    )

incLevel :: TermTypeM a -> TermTypeM a
incLevel :: forall a. TermTypeM a -> TermTypeM a
incLevel = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \TermEnv
env -> TermEnv
env {termLevel :: Int
termLevel = TermEnv -> Int
termLevel TermEnv
env forall a. Num a => a -> a -> a
+ Int
1}

-- | Get the type of an expression, with top level type variables
-- substituted.  Never call 'typeOf' directly (except in a few
-- carefully inspected locations)!
expType :: Exp -> TermTypeM StructType
expType :: Exp -> TermTypeM StructType
expType = forall (m :: * -> *). MonadUnify m => StructType -> m StructType
normType forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> StructType
typeOf

-- | Get the type of an expression, with all type variables
-- substituted.  Slower than 'expType', but sometimes necessary.
-- Never call 'typeOf' directly (except in a few carefully inspected
-- locations)!
expTypeFully :: Exp -> TermTypeM StructType
expTypeFully :: Exp -> TermTypeM StructType
expTypeFully = forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> StructType
typeOf

newArrayType :: Usage -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType :: Usage -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType Usage
usage Name
desc Int
r = do
  VName
v <- forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
desc
  VName -> Constraint -> TermTypeM ()
constrain VName
v forall a b. (a -> b) -> a -> b
$ Liftedness -> Usage -> Constraint
NoConstraint Liftedness
Unlifted Usage
usage
  [VName]
dims <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
r forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadUnify m =>
Usage -> Rigidity -> Name -> m VName
newDimVar Usage
usage Rigidity
Nonrigid Name
"dim"
  let rowt :: ScalarTypeBase dim NoUniqueness
rowt = forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar forall a. Monoid a => a
mempty (forall v. v -> QualName v
qualName VName
v) []
      mkSize :: VName -> Exp
mkSize = forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName (forall a. Located a => a -> SrcLoc
srclocOf Usage
usage) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName
  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array forall a. Monoid a => a
mempty (forall dim. [dim] -> Shape dim
Shape forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> Exp
mkSize [VName]
dims) forall {dim}. ScalarTypeBase dim NoUniqueness
rowt,
      forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall {dim}. ScalarTypeBase dim NoUniqueness
rowt
    )

-- | Replace *all* dimensions with distinct fresh size variables.
allDimsFreshInType ::
  Usage ->
  Rigidity ->
  Name ->
  TypeBase Size als ->
  TermTypeM (TypeBase Size als, M.Map VName Size)
allDimsFreshInType :: forall als.
Usage
-> Rigidity
-> Name
-> TypeBase Exp als
-> TermTypeM (TypeBase Exp als, Map VName Exp)
allDimsFreshInType Usage
usage Rigidity
r Name
desc TypeBase Exp als
t =
  forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse forall {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadTrans t, MonadUnify m, MonadState (Map VName a) (t m)) =>
a -> t m Exp
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase Exp als
t) forall a. Monoid a => a
mempty
  where
    onDim :: a -> t m Exp
onDim a
d = do
      VName
v <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadUnify m =>
Usage -> Rigidity -> Name -> m VName
newDimVar Usage
usage Rigidity
r Name
desc
      forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v a
d
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (forall v. v -> QualName v
qualName VName
v) forall a b. (a -> b) -> a -> b
$ forall a. Located a => a -> SrcLoc
srclocOf Usage
usage

-- | Replace all type variables with their concrete types.
updateTypes :: ASTMappable e => e -> TermTypeM e
updateTypes :: forall e. ASTMappable e => e -> TermTypeM e
updateTypes = forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper TermTypeM
tv
  where
    tv :: ASTMapper TermTypeM
tv =
      ASTMapper
        { mapOnExp :: Exp -> TermTypeM Exp
mapOnExp = forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper TermTypeM
tv,
          mapOnName :: VName -> TermTypeM VName
mapOnName = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnStructType :: StructType -> TermTypeM StructType
mapOnStructType = forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully,
          mapOnParamType :: ParamType -> TermTypeM ParamType
mapOnParamType = forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully,
          mapOnResRetType :: RetTypeBase Exp Uniqueness
-> TermTypeM (RetTypeBase Exp Uniqueness)
mapOnResRetType = forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully
        }

--- Basic checking

unifies :: T.Text -> StructType -> Exp -> TermTypeM Exp
unifies :: Text -> StructType -> Exp -> TermTypeM Exp
unifies Text
why StructType
t Exp
e = do
  forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (forall a. Located a => a -> Text -> Usage
mkUsage (forall a. Located a => a -> SrcLoc
srclocOf Exp
e) Text
why) StructType
t forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM StructType
expType Exp
e
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

-- | @require ts e@ causes a 'TypeError' if @expType e@ is not one of
-- the types in @ts@.  Otherwise, simply returns @e@.
require :: T.Text -> [PrimType] -> Exp -> TermTypeM Exp
require :: Text -> [PrimType] -> Exp -> TermTypeM Exp
require Text
why [PrimType]
ts Exp
e = do
  forall (m :: * -> *).
MonadUnify m =>
[PrimType] -> Usage -> StructType -> m ()
mustBeOneOf [PrimType]
ts (forall a. Located a => a -> Text -> Usage
mkUsage (forall a. Located a => a -> SrcLoc
srclocOf Exp
e) Text
why) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> TermTypeM StructType
expType Exp
e
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

termCheckTypeExp ::
  TypeExp NoInfo Name ->
  TermTypeM (TypeExp Info VName, [VName], ResRetType)
termCheckTypeExp :: TypeExp NoInfo Name
-> TermTypeM
     (TypeExp Info VName, [VName], RetTypeBase Exp Uniqueness)
termCheckTypeExp TypeExp NoInfo Name
te = do
  (TypeExp Info VName
te', [VName]
svars, RetTypeBase Exp Uniqueness
rettype, Liftedness
_l) <- forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp NoInfo Name
-> m (TypeExp Info VName, [VName], RetTypeBase Exp Uniqueness,
      Liftedness)
checkTypeExp TypeExp NoInfo Name
te

  -- No guarantee that the locally bound sizes in rettype are globally
  -- unique, but we want to turn them into size variables, so let's
  -- give them some unique names.  Maybe this should be done below,
  -- where we actually turn these into size variables?
  RetType [VName]
dims ResType
st <- forall (m :: * -> *).
MonadTypeChecker m =>
RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness)
renameRetType RetTypeBase Exp Uniqueness
rettype

  forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp Info VName
te', [VName]
svars, forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims ResType
st)

checkTypeExpNonrigid :: TypeExp NoInfo Name -> TermTypeM (TypeExp Info VName, ResType, [VName])
checkTypeExpNonrigid :: TypeExp NoInfo Name
-> TermTypeM (TypeExp Info VName, ResType, [VName])
checkTypeExpNonrigid TypeExp NoInfo Name
te = do
  (TypeExp Info VName
te', [VName]
svars, RetType [VName]
dims ResType
st) <- TypeExp NoInfo Name
-> TermTypeM
     (TypeExp Info VName, [VName], RetTypeBase Exp Uniqueness)
termCheckTypeExp TypeExp NoInfo Name
te
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName]
svars forall a. [a] -> [a] -> [a]
++ [VName]
dims) forall a b. (a -> b) -> a -> b
$ \VName
v ->
    VName -> Constraint -> TermTypeM ()
constrain VName
v forall a b. (a -> b) -> a -> b
$ Maybe Exp -> Usage -> Constraint
Size forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ forall a. Located a => a -> Text -> Usage
mkUsage (forall a. Located a => a -> SrcLoc
srclocOf TypeExp NoInfo Name
te) Text
"anonymous size in type expression"
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp Info VName
te', ResType
st, [VName]
svars forall a. [a] -> [a] -> [a]
++ [VName]
dims)

--- Sizes

isInt64 :: Exp -> Maybe Int64
isInt64 :: Exp -> Maybe Int64
isInt64 (Literal (SignedValue (Int64Value Int64
k')) SrcLoc
_) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
k'
isInt64 (IntLit Integer
k' Info StructType
_ SrcLoc
_) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger Integer
k'
isInt64 (Negate Exp
x SrcLoc
_) = forall a. Num a => a -> a
negate forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Maybe Int64
isInt64 Exp
x
isInt64 (Parens Exp
x SrcLoc
_) = Exp -> Maybe Int64
isInt64 Exp
x
isInt64 Exp
_ = forall a. Maybe a
Nothing

-- Running

initialTermScope :: TermScope
initialTermScope :: TermScope
initialTermScope =
  TermScope
    { scopeVtable :: Map VName ValBinding
scopeVtable = Map VName ValBinding
initialVtable,
      scopeTypeTable :: Map VName TypeBinding
scopeTypeTable = forall a. Monoid a => a
mempty,
      scopeNameMap :: NameMap
scopeNameMap = forall a. Monoid a => a
mempty,
      scopeModTable :: Map VName Mod
scopeModTable = forall a. Monoid a => a
mempty
    }
  where
    initialVtable :: Map VName ValBinding
initialVtable = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe forall {a}. (a, Intrinsic) -> Maybe (a, ValBinding)
addIntrinsicF forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map VName Intrinsic
intrinsics

    prim :: PrimType -> TypeBase dim u
prim = forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u. PrimType -> ScalarTypeBase dim u
Prim
    arrow :: TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness -> TypeBase dim u
arrow TypeBase dim NoUniqueness
x RetTypeBase dim Uniqueness
y = forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe TypeBase dim NoUniqueness
x RetTypeBase dim Uniqueness
y

    addIntrinsicF :: (a, Intrinsic) -> Maybe (a, ValBinding)
addIntrinsicF (a
name, IntrinsicMonoFun [PrimType]
pts PrimType
t) =
      forall a. a -> Maybe a
Just (a
name, [TypeParam] -> StructType -> ValBinding
BoundV [] forall a b. (a -> b) -> a -> b
$ forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness -> TypeBase dim u
arrow forall {dim} {u}. TypeBase dim u
pts' forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall {dim} {u}. PrimType -> TypeBase dim u
prim PrimType
t)
      where
        pts' :: TypeBase dim u
pts' = case [PrimType]
pts of
          [PrimType
pt] -> forall {dim} {u}. PrimType -> TypeBase dim u
prim PrimType
pt
          [PrimType]
_ -> forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall {dim} {u}. PrimType -> TypeBase dim u
prim [PrimType]
pts
    addIntrinsicF (a
name, IntrinsicOverloadedFun [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rts) =
      forall a. a -> Maybe a
Just (a
name, [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> ValBinding
OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rts)
    addIntrinsicF (a
name, IntrinsicPolyFun [TypeParam]
tvs [ParamType]
pts RetTypeBase Exp Uniqueness
rt) =
      forall a. a -> Maybe a
Just
        ( a
name,
          [TypeParam] -> StructType -> ValBinding
BoundV [TypeParam]
tvs forall a b. (a -> b) -> a -> b
$ [ParamType] -> RetTypeBase Exp Uniqueness -> StructType
foldFunType [ParamType]
pts RetTypeBase Exp Uniqueness
rt
        )
    addIntrinsicF (a
name, Intrinsic
IntrinsicEquality) =
      forall a. a -> Maybe a
Just (a
name, ValBinding
EqualityF)
    addIntrinsicF (a, Intrinsic)
_ = forall a. Maybe a
Nothing

runTermTypeM :: (UncheckedExp -> TermTypeM Exp) -> TermTypeM a -> TypeM a
runTermTypeM :: forall a. (UncheckedExp -> TermTypeM Exp) -> TermTypeM a -> TypeM a
runTermTypeM UncheckedExp -> TermTypeM Exp
checker (TermTypeM ReaderT
  TermEnv (StateT TermTypeState (Except (Warnings, TypeError))) a
m) = do
  TermScope
initial_scope <- (TermScope
initialTermScope <>) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> TermScope
envToTermScope forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeM Env
askEnv
  ImportName
name <- TypeM ImportName
askImportName
  Env
outer_env <- TypeM Env
askEnv
  VNameSource
src <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TypeState -> VNameSource
TypeM.stateNameSource
  let initial_tenv :: TermEnv
initial_tenv =
        TermEnv
          { termScope :: TermScope
termScope = TermScope
initial_scope,
            termChecking :: Maybe Checking
termChecking = forall a. Maybe a
Nothing,
            termLevel :: Int
termLevel = Int
0,
            termChecker :: UncheckedExp -> TermTypeM Exp
termChecker = UncheckedExp -> TermTypeM Exp
checker,
            termImportName :: ImportName
termImportName = ImportName
name,
            termOuterEnv :: Env
termOuterEnv = Env
outer_env
          }
      initial_state :: TermTypeState
initial_state =
        TermTypeState
          { stateConstraints :: Constraints
stateConstraints = forall a. Monoid a => a
mempty,
            stateCounter :: Int
stateCounter = Int
0,
            stateUsed :: Set VName
stateUsed = forall a. Monoid a => a
mempty,
            stateWarnings :: Warnings
stateWarnings = forall a. Monoid a => a
mempty,
            stateNameSource :: VNameSource
stateNameSource = VNameSource
src
          }
  case forall e a. Except e a -> Either e a
runExcept (forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT
  TermEnv (StateT TermTypeState (Except (Warnings, TypeError))) a
m TermEnv
initial_tenv) TermTypeState
initial_state) of
    Left (Warnings
ws, TypeError
e) -> do
      forall (m :: * -> *). MonadTypeChecker m => Warnings -> m ()
warnings Warnings
ws
      forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError TypeError
e
    Right (a
a, TermTypeState {VNameSource
stateNameSource :: VNameSource
stateNameSource :: TermTypeState -> VNameSource
stateNameSource, Warnings
stateWarnings :: Warnings
stateWarnings :: TermTypeState -> Warnings
stateWarnings}) -> do
      forall (m :: * -> *). MonadTypeChecker m => Warnings -> m ()
warnings Warnings
stateWarnings
      forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \TypeState
s -> TypeState
s {stateNameSource :: VNameSource
TypeM.stateNameSource = VNameSource
stateNameSource}
      forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a