{-# LANGUAGE Strict #-}

-- | Unsized type checking.
--
-- This checker generates type constraints (type 'CtTy') which are then solved
-- to find a solution. The result is a decorated AST where most of the type
-- annotations are just references to type variables. Further, all the
-- size-specific annotations (e.g. existential sizes) just contain dummy values,
-- such as empty lists.
--
-- If Futhark had no fancy type system features, then this pass would
-- essentially be all you needed.
module Language.Futhark.TypeChecker.Terms.Unsized
  ( checkValDef,
    checkSingleExp,
    checkSizeExp,
    Solution,
  )
where

import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Bifoldable (bifoldMap)
import Data.Bifunctor
import Data.Bitraversable
import Data.Char (isAscii)
import Data.Either (partitionEithers)
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Loc (Loc (NoLoc))
import Data.Map qualified as M
import Data.Maybe
import Data.Ord (comparing)
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.FreshNames qualified as FreshNames
import Futhark.MonadFreshNames hiding (newName)
import Futhark.Util (nubOrd)
import Futhark.Util.Pretty
import Language.Futhark
import Language.Futhark.TypeChecker.Constraints
import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupMod)
import Language.Futhark.TypeChecker.Monad qualified as TypeM
import Language.Futhark.TypeChecker.Terms.Scope hiding (envToTermScope, initialTermScope, lookupQualNameEnv)
import Language.Futhark.TypeChecker.Terms.Scope qualified as Scope
import Language.Futhark.TypeChecker.TySolve hiding (Type)
import Language.Futhark.TypeChecker.Types
import Language.Futhark.TypeChecker.Unify (mkUsage)
import Prelude hiding (mod)

type Type = CtType ()

-- | The unsized checker ignores sizes entirely, so we erase them. (The
-- constraint solver never inspects sizes; see '()' for the shape
-- representation this could use if size-aware rank inference were ever
-- wired up.)
toType :: TypeBase Size o -> TypeBase () o
toType :: forall o. TypeBase Exp o -> TypeBase () o
toType = (Exp -> ()) -> TypeBase Exp o -> TypeBase () o
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> Exp -> ()
forall a b. a -> b -> a
const ())

-- | 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 -> Int
termLevel :: Level,
    TermEnv -> Env
termOuterEnv :: Env,
    TermEnv -> ImportName
termImportName :: ImportName,
    -- | The liftedness of abstract types.
    TermEnv -> TySet
termTySet :: TySet
  }

-- | An instantiation (at the given location, of a type parameter of
-- the given polymorphic function, with the given liftedness, as the
-- given type variable), recorded such that we can check, after
-- constraint solving, that the liftedness of the type parameter is
-- respected ('checkTyInstLiftedness').
data TyInst = TyInst Loc (QualName VName) Liftedness TyVar

-- | 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 TermState = TermState
  { TermState -> [CtTy ()]
termConstraints :: [CtTy ()],
    TermState -> TyVars ()
termTyVars :: TyVars (),
    TermState -> TyParams
termTyParams :: TyParams,
    TermState -> [TyInst]
termTyInsts :: [TyInst],
    TermState -> Int
termCounter :: !Int,
    TermState -> Warnings
termWarnings :: Warnings,
    TermState -> VNameSource
termNameSource :: VNameSource,
    -- | Mapping from artificial type variables to the actual types they represent.
    TermState -> Map TyVar (TypeBase () NoMode)
termArtificial :: M.Map TyVar Type
  }

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

-- | The scope, with sizes erased. See
-- "Language.Futhark.TypeChecker.Terms.Scope".
envToTermScope :: Env -> TermScope ()
envToTermScope :: Env -> TermScope ()
envToTermScope = (StructType -> TypeBase () NoMode) -> Env -> TermScope ()
forall dim.
(StructType -> TypeBase dim NoMode) -> Env -> TermScope dim
Scope.envToTermScope StructType -> TypeBase () NoMode
forall o. TypeBase Exp o -> TypeBase () o
toType

initialTermScope :: TermScope ()
initialTermScope :: TermScope ()
initialTermScope = (StructType -> TypeBase () NoMode) -> TermScope ()
forall dim. (StructType -> TypeBase dim NoMode) -> TermScope dim
Scope.initialTermScope StructType -> TypeBase () NoMode
forall o. TypeBase Exp o -> TypeBase () o
toType

runTermM :: TermM a -> TypeM a
runTermM :: forall a. TermM a -> TypeM a
runTermM (TermM ReaderT TermEnv (StateT TermState (Except (Warnings, TypeError))) a
m) = do
  initial_scope <- (TermScope ()
initialTermScope TermScope () -> TermScope () -> TermScope ()
forall a. Semigroup a => a -> a -> a
<>) (TermScope () -> TermScope ())
-> (Env -> TermScope ()) -> Env -> TermScope ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> TermScope ()
forall dim. Env -> TermScope dim
Scope.envToTermScopeNoVals (Env -> TermScope ()) -> TypeM Env -> TypeM (TermScope ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeM Env
askEnv
  name <- askImportName
  outer_env <- askEnv
  src <- gets stateNameSource
  abs_types <- getTySet
  let initial_env =
        TermEnv
          { termScope :: TermScope ()
termScope = TermScope ()
initial_scope,
            termLevel :: Int
termLevel = Int
0,
            termImportName :: ImportName
termImportName = ImportName
name,
            termOuterEnv :: Env
termOuterEnv = Env
outer_env,
            termTySet :: TySet
termTySet = TySet
abs_types
          }
      initial_state =
        TermState
          { termConstraints :: [CtTy ()]
termConstraints = [CtTy ()]
forall a. Monoid a => a
mempty,
            termTyVars :: TyVars ()
termTyVars = TyVars ()
forall a. Monoid a => a
mempty,
            termTyParams :: TyParams
termTyParams = TyParams
forall a. Monoid a => a
mempty,
            termTyInsts :: [TyInst]
termTyInsts = [TyInst]
forall a. Monoid a => a
mempty,
            termWarnings :: Warnings
termWarnings = Warnings
forall a. Monoid a => a
mempty,
            termNameSource :: VNameSource
termNameSource = VNameSource
src,
            termCounter :: Int
termCounter = Int
0,
            termArtificial :: Map TyVar (TypeBase () NoMode)
termArtificial = Map TyVar (TypeBase () NoMode)
forall a. Monoid a => a
mempty
          }
  case runExcept (runStateT (runReaderT m initial_env) initial_state) of
    Left (Warnings
ws, TypeError
e) -> do
      Warnings -> TypeM ()
forall (m :: * -> *). MonadTypeChecker m => Warnings -> m ()
warnings Warnings
ws
      TypeError -> TypeM a
forall a. TypeError -> TypeM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError TypeError
e
    Right (a
a, TermState {VNameSource
termNameSource :: TermState -> VNameSource
termNameSource :: VNameSource
termNameSource, Warnings
termWarnings :: TermState -> Warnings
termWarnings :: Warnings
termWarnings}) -> do
      Warnings -> TypeM ()
forall (m :: * -> *). MonadTypeChecker m => Warnings -> m ()
warnings Warnings
termWarnings
      (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 {stateNameSource = termNameSource}
      a -> TypeM a
forall a. a -> TypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

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

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

incCounter :: TermM Int
incCounter :: TermM Int
incCounter = do
  s <- TermM TermState
forall s (m :: * -> *). MonadState s m => m s
get
  put s {termCounter = termCounter s + 1}
  pure $ termCounter s

tyVarType :: o -> TyVar -> TypeBase dim o
tyVarType :: forall o dim. o -> TyVar -> TypeBase dim o
tyVarType o
o TyVar
v = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ o -> QualName TyVar -> [TypeArg dim] -> ScalarTypeBase dim o
forall dim o.
o -> QualName TyVar -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o (TyVar -> QualName TyVar
forall v. v -> QualName v
qualName TyVar
v) []

newTyVarWith :: Name -> TyVarInfo () -> TermM TyVar
newTyVarWith :: Name -> TyVarInfo () -> TermM TyVar
newTyVarWith Name
desc TyVarInfo ()
info = do
  i <- TermM Int
incCounter
  v <- newID $ mkTypeVarName desc i
  lvl <- curLevel
  modify $ \TermState
s -> TermState
s {termTyVars = M.insert v (lvl, info) $ termTyVars s}
  pure v

newTyVar :: (Located loc) => loc -> Liftedness -> Name -> TermM TyVar
newTyVar :: forall loc. Located loc => loc -> Liftedness -> Name -> TermM TyVar
newTyVar loc
loc Liftedness
l Name
desc = Name -> TyVarInfo () -> TermM TyVar
newTyVarWith Name
desc (TyVarInfo () -> TermM TyVar) -> TyVarInfo () -> TermM TyVar
forall a b. (a -> b) -> a -> b
$ Loc -> Liftedness -> TyVarInfo ()
forall d. Loc -> Liftedness -> TyVarInfo d
TyVarFree (loc -> Loc
forall a. Located a => a -> Loc
locOf loc
loc) Liftedness
l

newType :: (Located loc) => loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType :: forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType loc
loc Liftedness
l Name
desc o
o = o -> TyVar -> TypeBase dim o
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType o
o (TyVar -> TypeBase dim o) -> TermM TyVar -> TermM (TypeBase dim o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> loc -> Liftedness -> Name -> TermM TyVar
forall loc. Located loc => loc -> Liftedness -> Name -> TermM TyVar
newTyVar loc
loc Liftedness
l Name
desc

-- | New type that must be allowed as an array element.
newElemType :: (Located loc) => loc -> Name -> o -> TermM (TypeBase dim o)
newElemType :: forall loc o dim.
Located loc =>
loc -> Name -> o -> TermM (TypeBase dim o)
newElemType loc
loc Name
desc o
o = o -> TyVar -> TypeBase dim o
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType o
o (TyVar -> TypeBase dim o) -> TermM TyVar -> TermM (TypeBase dim o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> loc -> Liftedness -> Name -> TermM TyVar
forall loc. Located loc => loc -> Liftedness -> Name -> TermM TyVar
newTyVar loc
loc Liftedness
Unlifted Name
desc

newTypeWithField :: SrcLoc -> Name -> Name -> Type -> TermM Type
newTypeWithField :: SrcLoc
-> Name -> Name -> TypeBase () NoMode -> TermM (TypeBase () NoMode)
newTypeWithField SrcLoc
loc Name
desc Name
k TypeBase () NoMode
t =
  NoMode -> TyVar -> TypeBase () NoMode
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType NoMode
NoMode
    (TyVar -> TypeBase () NoMode)
-> TermM TyVar -> TermM (TypeBase () NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> TyVarInfo () -> TermM TyVar
newTyVarWith Name
desc (Loc -> Map Name (TypeBase () NoMode) -> TyVarInfo ()
forall d. Loc -> Map Name (CtType d) -> TyVarInfo d
TyVarRecord (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Map Name (TypeBase () NoMode) -> TyVarInfo ())
-> Map Name (TypeBase () NoMode) -> TyVarInfo ()
forall a b. (a -> b) -> a -> b
$ Name -> TypeBase () NoMode -> Map Name (TypeBase () NoMode)
forall k a. k -> a -> Map k a
M.singleton Name
k TypeBase () NoMode
t)

newTypeWithConstr :: SrcLoc -> Name -> o -> Name -> [TypeBase () o] -> TermM (TypeBase d o)
newTypeWithConstr :: forall o d.
SrcLoc
-> Name -> o -> Name -> [TypeBase () o] -> TermM (TypeBase d o)
newTypeWithConstr SrcLoc
loc Name
desc o
o Name
k [TypeBase () o]
ts =
  o -> TyVar -> TypeBase d o
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType o
o (TyVar -> TypeBase d o) -> TermM TyVar -> TermM (TypeBase d o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> TyVarInfo () -> TermM TyVar
newTyVarWith Name
desc (Loc -> Map Name [TypeBase () NoMode] -> TyVarInfo ()
forall d. Loc -> Map Name [CtType d] -> TyVarInfo d
TyVarSum (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Map Name [TypeBase () NoMode] -> TyVarInfo ())
-> Map Name [TypeBase () NoMode] -> TyVarInfo ()
forall a b. (a -> b) -> a -> b
$ Name -> [TypeBase () NoMode] -> Map Name [TypeBase () NoMode]
forall k a. k -> a -> Map k a
M.singleton Name
k [TypeBase () NoMode]
ts')
  where
    ts' :: [TypeBase () NoMode]
ts' = (TypeBase () o -> TypeBase () NoMode)
-> [TypeBase () o] -> [TypeBase () NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase () o -> NoMode -> TypeBase () NoMode
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`setMode` NoMode
NoMode) [TypeBase () o]
ts

newTypeOverloaded :: SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded :: forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded SrcLoc
loc Name
name [PrimType]
pts =
  NoMode -> TyVar -> TypeBase d NoMode
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType NoMode
NoMode (TyVar -> TypeBase d NoMode)
-> TermM TyVar -> TermM (TypeBase d NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> TyVarInfo () -> TermM TyVar
newTyVarWith Name
name (Loc -> [PrimType] -> TyVarInfo ()
forall d. Loc -> [PrimType] -> TyVarInfo d
TyVarPrim (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) [PrimType]
pts)

newArtificial :: o -> TypeBase () o -> TermM (TypeBase Size o)
newArtificial :: forall o. o -> TypeBase () o -> TermM (TypeBase Exp o)
newArtificial o
o TypeBase () o
t = do
  v <- Name -> TermM TyVar
forall (m :: * -> *). MonadTypeChecker m => Name -> m TyVar
newID Name
"artificial"
  let t' = o -> TyVar -> TypeBase dim o
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType o
o TyVar
v
  modify $ \TermState
s -> TermState
s {termArtificial = M.insert v (second (const NoMode) t) $ termArtificial s}
  pure t'

-- The AST requires annotations to be StructTypes, but the type
-- checker works with Types. This creates artificial type "variables"
-- that allow us to connect the AST annotations with the actual
-- inferred types. The artificial variables should never occur in
-- constraints - they can be substituted away with asType.
--
-- Equal components (with fully known shapes) of the same annotation
-- are given the same artificial variable, so that the sized type
-- checker knows that they have the same sizes.
asStructType :: TypeBase () o -> TermM (TypeBase Size o)
asStructType :: forall o. TypeBase () o -> TermM (TypeBase Exp o)
asStructType TypeBase () o
t = StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o)
-> Map (TypeBase () NoMode) TyVar -> TermM (TypeBase Exp o)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (TypeBase () o
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o)
forall o'.
TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
onType TypeBase () o
t) Map (TypeBase () NoMode) TyVar
forall a. Monoid a => a
mempty
  where
    onType ::
      TypeBase () o' ->
      StateT (M.Map (TypeBase () NoMode) TyVar) TermM (TypeBase Size o')
    onType :: forall o'.
TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
onType (Scalar (Prim PrimType
pt)) = TypeBase Exp o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall a. a -> StateT (Map (TypeBase () NoMode) TyVar) TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp o'
 -> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o'))
-> TypeBase Exp o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp o' -> TypeBase Exp o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o' -> TypeBase Exp o')
-> ScalarTypeBase Exp o' -> TypeBase Exp o'
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp o'
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
pt
    onType (Scalar (TypeVar o'
o QualName TyVar
v [])) = TypeBase Exp o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall a. a -> StateT (Map (TypeBase () NoMode) TyVar) TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp o'
 -> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o'))
-> TypeBase Exp o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp o' -> TypeBase Exp o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o' -> TypeBase Exp o')
-> ScalarTypeBase Exp o' -> TypeBase Exp o'
forall a b. (a -> b) -> a -> b
$ o' -> QualName TyVar -> [TypeArg Exp] -> ScalarTypeBase Exp o'
forall dim o.
o -> QualName TyVar -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o'
o QualName TyVar
v []
    onType (Scalar (Arrow o'
o PName
pname Diet
d TypeBase () NoMode
t1 (RetType [TyVar]
ext TypeBase () Freshness
t2))) = do
      t1' <- TypeBase () NoMode
-> StateT (Map (TypeBase () NoMode) TyVar) TermM StructType
forall o'.
TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
onType TypeBase () NoMode
t1
      t2' <- onType t2
      pure $ Scalar $ Arrow o pname d t1' $ RetType ext t2'
    onType (Scalar (Record Map Name (TypeBase () o')
fs)) =
      ScalarTypeBase Exp o' -> TypeBase Exp o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o' -> TypeBase Exp o')
-> (Map Name (TypeBase Exp o') -> ScalarTypeBase Exp o')
-> Map Name (TypeBase Exp o')
-> TypeBase Exp o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase Exp o') -> ScalarTypeBase Exp o'
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o') -> TypeBase Exp o')
-> StateT
     (Map (TypeBase () NoMode) TyVar) TermM (Map Name (TypeBase Exp o'))
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase () o'
 -> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o'))
-> Map Name (TypeBase () o')
-> StateT
     (Map (TypeBase () NoMode) TyVar) TermM (Map Name (TypeBase Exp o'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall o'.
TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
onType Map Name (TypeBase () o')
fs
    onType (Scalar (Sum Map Name [TypeBase () o']
cs)) =
      ScalarTypeBase Exp o' -> TypeBase Exp o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o' -> TypeBase Exp o')
-> (Map Name [TypeBase Exp o'] -> ScalarTypeBase Exp o')
-> Map Name [TypeBase Exp o']
-> TypeBase Exp o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase Exp o'] -> ScalarTypeBase Exp o'
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o'] -> TypeBase Exp o')
-> StateT
     (Map (TypeBase () NoMode) TyVar) TermM (Map Name [TypeBase Exp o'])
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase () o']
 -> StateT (Map (TypeBase () NoMode) TyVar) TermM [TypeBase Exp o'])
-> Map Name [TypeBase () o']
-> StateT
     (Map (TypeBase () NoMode) TyVar) TermM (Map Name [TypeBase Exp o'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse ((TypeBase () o'
 -> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o'))
-> [TypeBase () o']
-> StateT (Map (TypeBase () NoMode) TyVar) TermM [TypeBase Exp o']
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall o'.
TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
onType) Map Name [TypeBase () o']
cs
    onType t' :: TypeBase () o'
t'@(Scalar (TypeVar o'
o QualName TyVar
_ [TypeArg ()]
_)) = o'
-> TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall {t :: (* -> *) -> * -> *} {o}.
(MonadTrans t,
 MonadState (Map (TypeBase () NoMode) TyVar) (t TermM)) =>
o -> TypeBase () o -> t TermM (TypeBase Exp o)
artificial o'
o TypeBase () o'
t'
    onType t' :: TypeBase () o'
t'@(Array o'
o Shape ()
_ ScalarTypeBase () NoMode
_) = o'
-> TypeBase () o'
-> StateT (Map (TypeBase () NoMode) TyVar) TermM (TypeBase Exp o')
forall {t :: (* -> *) -> * -> *} {o}.
(MonadTrans t,
 MonadState (Map (TypeBase () NoMode) TyVar) (t TermM)) =>
o -> TypeBase () o -> t TermM (TypeBase Exp o)
artificial o'
o TypeBase () o'
t'

    artificial :: o -> TypeBase () o -> t TermM (TypeBase Exp o)
artificial o
o TypeBase () o
t'
      | TypeBase () o -> Bool
forall {b}. TypeBase () b -> Bool
anonymousShape TypeBase () o
t' = TermM (TypeBase Exp o) -> t TermM (TypeBase Exp o)
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (TermM (TypeBase Exp o) -> t TermM (TypeBase Exp o))
-> TermM (TypeBase Exp o) -> t TermM (TypeBase Exp o)
forall a b. (a -> b) -> a -> b
$ o -> TypeBase () o -> TermM (TypeBase Exp o)
forall o. o -> TypeBase () o -> TermM (TypeBase Exp o)
newArtificial o
o TypeBase () o
t'
      | Bool
otherwise = do
          let key :: TypeBase () NoMode
key = (o -> NoMode) -> TypeBase () o -> TypeBase () NoMode
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (NoMode -> o -> NoMode
forall a b. a -> b -> a
const NoMode
NoMode) TypeBase () o
t'
          seen <- t TermM (Map (TypeBase () NoMode) TyVar)
forall s (m :: * -> *). MonadState s m => m s
get
          case M.lookup key seen of
            Just TyVar
v -> TypeBase Exp o -> t TermM (TypeBase Exp o)
forall a. a -> t TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp o -> t TermM (TypeBase Exp o))
-> TypeBase Exp o -> t TermM (TypeBase Exp o)
forall a b. (a -> b) -> a -> b
$ o -> TyVar -> TypeBase Exp o
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType o
o TyVar
v
            Maybe TyVar
Nothing -> do
              v <- TermM TyVar -> t TermM TyVar
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (TermM TyVar -> t TermM TyVar) -> TermM TyVar -> t TermM TyVar
forall a b. (a -> b) -> a -> b
$ Name -> TermM TyVar
forall (m :: * -> *). MonadTypeChecker m => Name -> m TyVar
newID Name
"artificial"
              lift $ modify $ \TermState
s ->
                TermState
s {termArtificial = M.insert v key $ termArtificial s}
              modify $ M.insert key v
              pure $ tyVarType o v

    anonymousShape :: TypeBase () b -> Bool
anonymousShape = () -> [()] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem () ([()] -> Bool) -> (TypeBase () b -> [()]) -> TypeBase () b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (() -> [()]) -> (b -> [()]) -> TypeBase () b -> [()]
forall m a b. Monoid m => (a -> m) -> (b -> m) -> TypeBase a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap (() -> [()] -> [()]
forall a. a -> [a] -> [a]
: []) ([()] -> b -> [()]
forall a b. a -> b -> a
const [()]
forall a. Monoid a => a
mempty)

asType :: (Monoid o) => TypeBase Size o -> TermM (TypeBase () o)
asType :: forall o. Monoid o => TypeBase Exp o -> TermM (TypeBase () o)
asType TypeBase Exp o
t = do
  artificial <- (TermState -> Map TyVar (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TermState -> Map TyVar (TypeBase () NoMode)
termArtificial
  pure $ substTyVars (`M.lookup` artificial) (toType t)

expType :: Exp -> TermM Type
expType :: Exp -> TermM (TypeBase () NoMode)
expType = StructType -> TermM (TypeBase () NoMode)
forall o. Monoid o => TypeBase Exp o -> TermM (TypeBase () o)
asType (StructType -> TermM (TypeBase () NoMode))
-> (Exp -> StructType) -> Exp -> TermM (TypeBase () NoMode)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> StructType
typeOf -- NOTE: Only place you should use typeOf.

addCt :: CtTy () -> TermM ()
addCt :: CtTy () -> TermM ()
addCt CtTy ()
ct = (TermState -> TermState) -> TermM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TermState -> TermState) -> TermM ())
-> (TermState -> TermState) -> TermM ()
forall a b. (a -> b) -> a -> b
$ \TermState
s -> TermState
s {termConstraints = ct : termConstraints s}

ctEq :: Reason (CtType ()) -> TypeBase () o1 -> TypeBase () o2 -> TermM ()
ctEq :: forall o1 o2.
Reason (TypeBase () NoMode)
-> TypeBase () o1 -> TypeBase () o2 -> TermM ()
ctEq Reason (TypeBase () NoMode)
reason TypeBase () o1
t1 TypeBase () o2
t2 =
  -- As a minor optimisation, do not add constraint if the types are
  -- equal.
  Bool -> TermM () -> TermM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (TypeBase () NoMode
t1' TypeBase () NoMode -> TypeBase () NoMode -> Bool
forall a. Eq a => a -> a -> Bool
== TypeBase () NoMode
t2') (TermM () -> TermM ()) -> TermM () -> TermM ()
forall a b. (a -> b) -> a -> b
$ CtTy () -> TermM ()
addCt (CtTy () -> TermM ()) -> CtTy () -> TermM ()
forall a b. (a -> b) -> a -> b
$ Reason (TypeBase () NoMode)
-> TypeBase () NoMode -> TypeBase () NoMode -> CtTy ()
forall d. Reason (CtType d) -> CtType d -> CtType d -> CtTy d
CtEq Reason (TypeBase () NoMode)
reason TypeBase () NoMode
t1' TypeBase () NoMode
t2'
  where
    t1' :: TypeBase () NoMode
t1' = TypeBase () o1
t1 TypeBase () o1 -> NoMode -> TypeBase () NoMode
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`setMode` NoMode
NoMode
    t2' :: TypeBase () NoMode
t2' = TypeBase () o2
t2 TypeBase () o2 -> NoMode -> TypeBase () NoMode
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`setMode` NoMode
NoMode

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

withEnv :: TermEnv -> Env -> TermEnv
withEnv :: TermEnv -> Env -> TermEnv
withEnv TermEnv
tenv Env
env = TermEnv
tenv {termScope = termScope tenv <> envToTermScope env}

lookupQualNameEnv :: QualName VName -> TermM (TermScope ())
lookupQualNameEnv :: QualName TyVar -> TermM (TermScope ())
lookupQualNameEnv QualName TyVar
qn = (TermEnv -> TermScope ()) -> TermM (TermScope ())
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((TermEnv -> TermScope ()) -> TermM (TermScope ()))
-> (TermEnv -> TermScope ()) -> TermM (TermScope ())
forall a b. (a -> b) -> a -> b
$ \TermEnv
tenv -> (StructType -> TypeBase () NoMode)
-> TermScope () -> QualName TyVar -> TermScope ()
forall dim.
(StructType -> TypeBase dim NoMode)
-> TermScope dim -> QualName TyVar -> TermScope dim
Scope.lookupQualNameEnv StructType -> TypeBase () NoMode
forall o. TypeBase Exp o -> TypeBase () o
toType (TermEnv -> TermScope ()
termScope TermEnv
tenv) QualName TyVar
qn

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

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

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

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

  newName :: TyVar -> TermM TyVar
newName TyVar
v = do
    s <- TermM TermState
forall s (m :: * -> *). MonadState s m => m s
get
    let (v', src') = FreshNames.newName (termNameSource s) v
    put $ s {termNameSource = src'}
    pure v'

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

  newTypeName :: Name -> TermM TyVar
newTypeName Name
name = do
    i <- TermM Int
incCounter
    newID $ mkTypeVarName name i

  bindVal :: forall a. TyVar -> BoundV -> TermM a -> TermM a
bindVal TyVar
v (TypeM.BoundV [TypeParam]
tps StructType
t) TermM a
m = do
    t' <- StructType -> TermM (TypeBase () NoMode)
forall o. Monoid o => TypeBase Exp o -> TermM (TypeBase () o)
asType StructType
t
    let f TermScope ()
scope = TermScope ()
scope {scopeVtable = M.insert v (BoundV tps t') $ scopeVtable scope}
    localScope f m

  lookupType :: QualName TyVar -> TermM ([TypeParam], StructRetType, Liftedness)
lookupType QualName TyVar
qn = do
    outer_env <- (TermEnv -> Env) -> TermM Env
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks TermEnv -> Env
termOuterEnv
    scope <- lookupQualNameEnv qn
    case M.lookup (qualLeaf qn) $ scopeTypeTable scope of
      Maybe TypeBinding
Nothing -> [Char] -> TermM ([TypeParam], StructRetType, Liftedness)
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermM ([TypeParam], StructRetType, Liftedness))
-> [Char] -> TermM ([TypeParam], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$ [Char]
"lookupType: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> QualName TyVar -> [Char]
forall a. Show a => a -> [Char]
show QualName TyVar
qn
      Just (TypeAbbr Liftedness
l [TypeParam]
ps (RetType [TyVar]
dims StructType
def)) ->
        ([TypeParam], StructRetType, Liftedness)
-> TermM ([TypeParam], StructRetType, Liftedness)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          ( [TypeParam]
ps,
            [TyVar] -> StructType -> StructRetType
forall dim o. [TyVar] -> TypeBase dim o -> RetTypeBase dim o
RetType [TyVar]
dims (StructType -> StructRetType) -> StructType -> StructRetType
forall a b. (a -> b) -> a -> b
$ Env -> [TyVar] -> [TyVar] -> StructType -> StructType
forall o.
Env -> [TyVar] -> [TyVar] -> TypeBase Exp o -> TypeBase Exp o
qualifyTypeVars Env
outer_env ((TypeParam -> TyVar) -> [TypeParam] -> [TyVar]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> TyVar
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
ps) (QualName TyVar -> [TyVar]
forall vn. QualName vn -> [vn]
qualQuals QualName TyVar
qn) StructType
def,
            Liftedness
l
          )

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

--- All the general machinery goes above.

arrayOfRank :: Int -> Type -> Type
arrayOfRank :: Int -> TypeBase () NoMode -> TypeBase () NoMode
arrayOfRank Int
n = Shape () -> TypeBase () NoMode -> TypeBase () NoMode
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf (Shape () -> TypeBase () NoMode -> TypeBase () NoMode)
-> Shape () -> TypeBase () NoMode -> TypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ [()] -> Shape ()
forall dim. [dim] -> Shape dim
Shape ([()] -> Shape ()) -> [()] -> Shape ()
forall a b. (a -> b) -> a -> b
$ Int -> () -> [()]
forall a. Int -> a -> [a]
replicate Int
n ()

require :: T.Text -> [PrimType] -> Exp -> TermM Exp
require :: Text -> [PrimType] -> Exp -> TermM Exp
require Text
_why [PrimType
pt] Exp
e = do
  e_t <- Exp -> TermM (TypeBase () NoMode)
expType Exp
e
  ctEq (Reason (locOf e)) (Scalar $ Prim pt) e_t
  pure e
require Text
_why [PrimType]
pts Exp
e = do
  t :: Type <- SrcLoc -> Name -> [PrimType] -> TermM (TypeBase () NoMode)
forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e) Name
"t" [PrimType]
pts
  e_t <- expType e
  ctEq (Reason (locOf e)) t e_t
  pure e

-- | 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.
instTypeScheme ::
  QualName VName ->
  SrcLoc ->
  [TypeParam] ->
  Type ->
  TermM ([VName], Type)
instTypeScheme :: QualName TyVar
-> SrcLoc
-> [TypeParam]
-> TypeBase () NoMode
-> TermM ([TyVar], TypeBase () NoMode)
instTypeScheme QualName TyVar
qn SrcLoc
loc [TypeParam]
tparams TypeBase () NoMode
t = do
  (names, substs) <- ([Maybe (TyVar, (TyVar, TypeBase () NoMode))]
 -> ([TyVar], [(TyVar, TypeBase () NoMode)]))
-> TermM [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
-> TermM ([TyVar], [(TyVar, TypeBase () NoMode)])
forall a b. (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(TyVar, (TyVar, TypeBase () NoMode))]
-> ([TyVar], [(TyVar, TypeBase () NoMode)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(TyVar, (TyVar, TypeBase () NoMode))]
 -> ([TyVar], [(TyVar, TypeBase () NoMode)]))
-> ([Maybe (TyVar, (TyVar, TypeBase () NoMode))]
    -> [(TyVar, (TyVar, TypeBase () NoMode))])
-> [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
-> ([TyVar], [(TyVar, TypeBase () NoMode)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
-> [(TyVar, (TyVar, TypeBase () NoMode))]
forall a. [Maybe a] -> [a]
catMaybes) (TermM [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
 -> TermM ([TyVar], [(TyVar, TypeBase () NoMode)]))
-> TermM [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
-> TermM ([TyVar], [(TyVar, TypeBase () NoMode)])
forall a b. (a -> b) -> a -> b
$
    [TypeParam]
-> (TypeParam
    -> TermM (Maybe (TyVar, (TyVar, TypeBase () NoMode))))
-> TermM [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [TypeParam]
tparams ((TypeParam -> TermM (Maybe (TyVar, (TyVar, TypeBase () NoMode))))
 -> TermM [Maybe (TyVar, (TyVar, TypeBase () NoMode))])
-> (TypeParam
    -> TermM (Maybe (TyVar, (TyVar, TypeBase () NoMode))))
-> TermM [Maybe (TyVar, (TyVar, TypeBase () NoMode))]
forall a b. (a -> b) -> a -> b
$ \TypeParam
tparam ->
      case TypeParam
tparam of
        TypeParamType Liftedness
l TyVar
v SrcLoc
_ -> do
          v' <- SrcLoc -> Liftedness -> Name -> TermM TyVar
forall loc. Located loc => loc -> Liftedness -> Name -> TermM TyVar
newTyVar SrcLoc
loc Liftedness
l (Name -> TermM TyVar) -> Name -> TermM TyVar
forall a b. (a -> b) -> a -> b
$ Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.takeWhile Char -> Bool
isAscii (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Name -> Text
nameToText (Name -> Text) -> Name -> Text
forall a b. (a -> b) -> a -> b
$ TyVar -> Name
baseName TyVar
v
          modify $
            \TermState
s -> TermState
s {termTyInsts = TyInst (locOf loc) qn l v' : termTyInsts s}
          pure $ Just (v, (typeParamName tparam, tyVarType NoMode v'))
        TypeParamDim {} ->
          Maybe (TyVar, (TyVar, TypeBase () NoMode))
-> TermM (Maybe (TyVar, (TyVar, TypeBase () NoMode)))
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (TyVar, (TyVar, TypeBase () NoMode))
forall a. Maybe a
Nothing
  let t' = (TyVar -> Maybe (TypeBase () NoMode))
-> TypeBase () NoMode -> TypeBase () NoMode
forall o d.
Monoid o =>
(TyVar -> Maybe (TypeBase d NoMode))
-> TypeBase d o -> TypeBase d o
substTyVars (TyVar
-> [(TyVar, TypeBase () NoMode)] -> Maybe (TypeBase () NoMode)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(TyVar, TypeBase () NoMode)]
substs) TypeBase () NoMode
t
  pure (names, t')

lookupMod :: QualName VName -> TermM Mod
lookupMod :: QualName TyVar -> TermM Mod
lookupMod qn :: QualName TyVar
qn@(QualName [TyVar]
_ TyVar
name) = do
  scope <- QualName TyVar -> TermM (TermScope ())
lookupQualNameEnv QualName TyVar
qn
  case M.lookup name $ scopeModTable scope of
    Maybe Mod
Nothing -> [Char] -> TermM Mod
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermM Mod) -> [Char] -> TermM Mod
forall a b. (a -> b) -> a -> b
$ [Char]
"lookupMod: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> QualName TyVar -> [Char]
forall a. Show a => a -> [Char]
show QualName TyVar
qn
    Just Mod
m -> Mod -> TermM Mod
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mod
m

lookupVar :: SrcLoc -> QualName VName -> TermM Type
lookupVar :: SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
loc qn :: QualName TyVar
qn@(QualName [TyVar]
qs TyVar
name) = do
  scope <- QualName TyVar -> TermM (TermScope ())
lookupQualNameEnv QualName TyVar
qn
  outer_env <- asks termOuterEnv
  -- Top-level value bindings are not in the term scope (see
  -- 'envToTermScopeNoVals'); look them up on demand in the outer
  -- environment.
  case M.lookup name (scopeVtable scope) `mplus` Scope.lookupOuterVal toType outer_env name of
    Maybe (ValBinding ())
Nothing ->
      [Char] -> TermM (TypeBase () NoMode)
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermM (TypeBase () NoMode))
-> [Char] -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ [Char]
"lookupVar: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> QualName TyVar -> [Char]
forall a. Show a => a -> [Char]
show QualName TyVar
qn
    Just (BoundV [TypeParam]
tparams TypeBase () NoMode
t) -> do
      if [TypeParam] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParam]
tparams Bool -> Bool -> Bool
&& [TyVar] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TyVar]
qs
        then TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase () NoMode
t
        else do
          (tnames, t') <- QualName TyVar
-> SrcLoc
-> [TypeParam]
-> TypeBase () NoMode
-> TermM ([TyVar], TypeBase () NoMode)
instTypeScheme QualName TyVar
qn SrcLoc
loc [TypeParam]
tparams TypeBase () NoMode
t
          -- Qualify abstract types, so that e.g. mismatch errors
          -- mention them by how they were accessed. The sizes need
          -- no qualification, as they are not even present here.
          pure $ qualifyTypeVarsWith (\Set TyVar -> QualName TyVar -> QualName TyVar
_ Set TyVar
_ ()
d -> ()
d) outer_env tnames qs t'
    Just ValBinding ()
EqualityF -> do
      argtype <- Diet -> TyVar -> TypeBase () Diet
forall o dim. o -> TyVar -> TypeBase dim o
tyVarType Diet
Observe (TyVar -> TypeBase () Diet)
-> TermM TyVar -> TermM (TypeBase () Diet)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> TyVarInfo () -> TermM TyVar
newTyVarWith Name
"t" (Loc -> Liftedness -> TyVarInfo ()
forall d. Loc -> Liftedness -> TyVarInfo d
TyVarFree (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) Liftedness
Unlifted)
      pure $ foldFunType [argtype, argtype] $ RetType [] $ Scalar $ Prim Bool
    Just (OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt) -> do
      argtype <- SrcLoc -> Name -> [PrimType] -> TermM (TypeBase () NoMode)
forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded SrcLoc
loc Name
"t" [PrimType]
ts
      let (pts', rt') = instOverloaded argtype pts rt
      pure $ foldFunType (map (second $ const Observe) pts') $ RetType [] $ second (const Nonfresh) rt'
    -- The unsized checker binds recursive functions directly (see
    -- 'checkRecursive'), so it never produces a 'RecursiveV'.
    Just ValBinding ()
RecursiveV ->
      [Char] -> TermM (TypeBase () NoMode)
forall a. HasCallStack => [Char] -> a
error ([Char] -> TermM (TypeBase () NoMode))
-> [Char] -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ [Char]
"lookupVar: unexpected RecursiveV for " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> QualName TyVar -> [Char]
forall a. Show a => a -> [Char]
show QualName TyVar
qn
  where
    instOverloaded :: TypeBase dim o
-> [Maybe PrimType]
-> Maybe PrimType
-> ([TypeBase dim o], TypeBase dim o)
instOverloaded TypeBase dim o
argtype [Maybe PrimType]
pts Maybe PrimType
rt =
      ( (Maybe PrimType -> TypeBase dim o)
-> [Maybe PrimType] -> [TypeBase dim o]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase dim o
-> (PrimType -> TypeBase dim o) -> Maybe PrimType -> TypeBase dim o
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeBase dim o
argtype (ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> (PrimType -> ScalarTypeBase dim o) -> PrimType -> TypeBase dim o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim)) [Maybe PrimType]
pts,
        TypeBase dim o
-> (PrimType -> TypeBase dim o) -> Maybe PrimType -> TypeBase dim o
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeBase dim o
argtype (ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> (PrimType -> ScalarTypeBase dim o) -> PrimType -> TypeBase dim o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim) Maybe PrimType
rt
      )

bind ::
  [Ident StructType] ->
  TermM a ->
  TermM a
bind :: forall a. [IdentBase Info TyVar StructType] -> TermM a -> TermM a
bind [IdentBase Info TyVar StructType]
idents TermM a
m = do
  let names :: [TyVar]
names = (IdentBase Info TyVar StructType -> TyVar)
-> [IdentBase Info TyVar StructType] -> [TyVar]
forall a b. (a -> b) -> [a] -> [b]
map IdentBase Info TyVar StructType -> TyVar
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName [IdentBase Info TyVar StructType]
idents
  ts <- (IdentBase Info TyVar StructType -> TermM (TypeBase () NoMode))
-> [IdentBase Info TyVar StructType] -> TermM [TypeBase () NoMode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (StructType -> TermM (TypeBase () NoMode)
forall o. Monoid o => TypeBase Exp o -> TermM (TypeBase () o)
asType (StructType -> TermM (TypeBase () NoMode))
-> (IdentBase Info TyVar StructType -> StructType)
-> IdentBase Info TyVar StructType
-> TermM (TypeBase () NoMode)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> StructType)
-> (IdentBase Info TyVar StructType -> Info StructType)
-> IdentBase Info TyVar StructType
-> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentBase Info TyVar StructType -> Info StructType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType) [IdentBase Info TyVar StructType]
idents
  localScope (`bindVars` zip names ts) m
  where
    bindVars :: TermScope dim -> [(TyVar, TypeBase dim NoMode)] -> TermScope dim
bindVars = (TermScope dim -> (TyVar, TypeBase dim NoMode) -> TermScope dim)
-> TermScope dim -> [(TyVar, TypeBase dim NoMode)] -> TermScope dim
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TermScope dim -> (TyVar, TypeBase dim NoMode) -> TermScope dim
forall {dim}.
TermScope dim -> (TyVar, TypeBase dim NoMode) -> TermScope dim
bindVar

    bindVar :: TermScope dim -> (TyVar, TypeBase dim NoMode) -> TermScope dim
bindVar TermScope dim
scope (TyVar
name, TypeBase dim NoMode
t) =
      TermScope dim
scope
        { scopeVtable = M.insert name (BoundV [] t) $ scopeVtable scope
        }

-- All this complexity is just so we can handle un-suffixed numeric
-- literals in patterns.
patLitMkType :: PatLit -> SrcLoc -> TermM ParamType
patLitMkType :: PatLit -> SrcLoc -> TermM ParamType
patLitMkType (PatLitInt Integer
_) SrcLoc
loc =
  Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe (StructType -> ParamType) -> TermM StructType -> TermM ParamType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc -> Name -> [PrimType] -> TermM StructType
forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded SrcLoc
loc Name
"t" [PrimType]
anyNumberType
patLitMkType (PatLitFloat Double
_) SrcLoc
loc =
  Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe (StructType -> ParamType) -> TermM StructType -> TermM ParamType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc -> Name -> [PrimType] -> TermM StructType
forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded SrcLoc
loc Name
"t" [PrimType]
anyFloatType
patLitMkType (PatLitPrim PrimValue
v) SrcLoc
_ =
  ParamType -> TermM ParamType
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ParamType -> TermM ParamType) -> ParamType -> TermM ParamType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
v

checkSizeExp' :: ExpBase NoInfo VName -> TermM Exp
checkSizeExp' :: ExpBase NoInfo TyVar -> TermM Exp
checkSizeExp' ExpBase NoInfo TyVar
e = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  e_t <- expType e'
  ctEq (Reason (locOf e)) e_t (Scalar (Prim (Signed Int64)))
  pure e'

checkPat' ::
  PatBase NoInfo VName ParamType ->
  Inferred ParamType ->
  TermM (Pat ParamType)
checkPat' :: PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' (PatParens PatBase NoInfo TyVar ParamType
p SrcLoc
loc) Inferred ParamType
t =
  Pat ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
PatBase f vn t -> SrcLoc -> PatBase f vn t
PatParens (Pat ParamType -> SrcLoc -> Pat ParamType)
-> TermM (Pat ParamType) -> TermM (SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' PatBase NoInfo TyVar ParamType
p Inferred ParamType
t TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkPat' (PatAttr AttrInfo TyVar
attr PatBase NoInfo TyVar ParamType
p SrcLoc
loc) Inferred ParamType
t =
  AttrInfo TyVar -> Pat ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
AttrInfo vn -> PatBase f vn t -> SrcLoc -> PatBase f vn t
PatAttr (AttrInfo TyVar -> Pat ParamType -> SrcLoc -> Pat ParamType)
-> TermM (AttrInfo TyVar)
-> TermM (Pat ParamType -> SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AttrInfo TyVar -> TermM (AttrInfo TyVar)
forall (m :: * -> *).
MonadTypeChecker m =>
AttrInfo TyVar -> m (AttrInfo TyVar)
checkAttr AttrInfo TyVar
attr TermM (Pat ParamType -> SrcLoc -> Pat ParamType)
-> TermM (Pat ParamType) -> TermM (SrcLoc -> Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' PatBase NoInfo TyVar ParamType
p Inferred ParamType
t TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkPat' (Id TyVar
name NoInfo ParamType
NoInfo SrcLoc
loc) (Ascribed ParamType
t) =
  Pat ParamType -> TermM (Pat ParamType)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat ParamType -> TermM (Pat ParamType))
-> Pat ParamType -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$ TyVar -> Info ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id TyVar
name (ParamType -> Info ParamType
forall a. a -> Info a
Info ParamType
t) SrcLoc
loc
checkPat' (Id TyVar
name NoInfo ParamType
NoInfo SrcLoc
loc) Inferred ParamType
NoneInferred = do
  t <- SrcLoc -> Liftedness -> Name -> Diet -> TermM ParamType
forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType SrcLoc
loc Liftedness
Lifted Name
"t" Diet
Observe
  pure $ Id name (Info t) loc
checkPat' (Wildcard NoInfo ParamType
_ SrcLoc
loc) (Ascribed ParamType
t) = do
  Pat ParamType -> TermM (Pat ParamType)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat ParamType -> TermM (Pat ParamType))
-> Pat ParamType -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$ Info ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (ParamType -> Info ParamType
forall a. a -> Info a
Info ParamType
t) SrcLoc
loc
checkPat' (Wildcard NoInfo ParamType
NoInfo SrcLoc
loc) Inferred ParamType
NoneInferred = do
  t <- SrcLoc -> Liftedness -> Name -> Diet -> TermM ParamType
forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType SrcLoc
loc Liftedness
Lifted Name
"t" Diet
Observe
  pure $ Wildcard (Info t) loc
checkPat' p :: PatBase NoInfo TyVar ParamType
p@(TuplePat [PatBase NoInfo TyVar ParamType]
ps SrcLoc
loc) (Ascribed ParamType
t)
  | Just [ParamType]
ts <- ParamType -> Maybe [ParamType]
forall dim o. TypeBase dim o -> Maybe [TypeBase dim o]
isTupleRecord ParamType
t,
    [ParamType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ParamType]
ts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [PatBase NoInfo TyVar ParamType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PatBase NoInfo TyVar ParamType]
ps =
      [Pat ParamType] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat
        ([Pat ParamType] -> SrcLoc -> Pat ParamType)
-> TermM [Pat ParamType] -> TermM (SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase NoInfo TyVar ParamType
 -> Inferred ParamType -> TermM (Pat ParamType))
-> [PatBase NoInfo TyVar ParamType]
-> [Inferred ParamType]
-> TermM [Pat ParamType]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' [PatBase NoInfo TyVar ParamType]
ps ((ParamType -> Inferred ParamType)
-> [ParamType] -> [Inferred ParamType]
forall a b. (a -> b) -> [a] -> [b]
map ParamType -> Inferred ParamType
forall t. t -> Inferred t
Ascribed [ParamType]
ts)
        TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  | Bool
otherwise =
      SrcLoc -> Notes -> Doc () -> TermM (Pat ParamType)
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM 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 () -> TermM (Pat ParamType))
-> Doc () -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Pattern"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (PatBase NoInfo TyVar ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. PatBase NoInfo TyVar ParamType -> Doc ann
pretty PatBase NoInfo TyVar ParamType
p)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"cannot match ascribed type"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ParamType -> Doc ann
pretty ParamType
t)
checkPat' (TuplePat [PatBase NoInfo TyVar ParamType]
ps SrcLoc
loc) Inferred ParamType
NoneInferred =
  [Pat ParamType] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat ([Pat ParamType] -> SrcLoc -> Pat ParamType)
-> TermM [Pat ParamType] -> TermM (SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase NoInfo TyVar ParamType -> TermM (Pat ParamType))
-> [PatBase NoInfo TyVar ParamType] -> TermM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
`checkPat'` Inferred ParamType
forall t. Inferred t
NoneInferred) [PatBase NoInfo TyVar ParamType]
ps TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkPat' p :: PatBase NoInfo TyVar ParamType
p@(RecordPat [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs SrcLoc
loc) Inferred ParamType
_
  | Just (L Loc
floc Name
f, PatBase NoInfo TyVar ParamType
_) <- ((L Name, PatBase NoInfo TyVar ParamType) -> Bool)
-> [(L Name, PatBase NoInfo TyVar ParamType)]
-> Maybe (L Name, PatBase NoInfo TyVar ParamType)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
L.find ((Text
"_" Text -> Text -> Bool
`T.isPrefixOf`) (Text -> Bool)
-> ((L Name, PatBase NoInfo TyVar ParamType) -> Text)
-> (L Name, PatBase NoInfo TyVar ParamType)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Text
nameToText (Name -> Text)
-> ((L Name, PatBase NoInfo TyVar ParamType) -> Name)
-> (L Name, PatBase NoInfo TyVar ParamType)
-> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. L Name -> Name
forall a. L a -> a
unLoc (L Name -> Name)
-> ((L Name, PatBase NoInfo TyVar ParamType) -> L Name)
-> (L Name, PatBase NoInfo TyVar ParamType)
-> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (L Name, PatBase NoInfo TyVar ParamType) -> L Name
forall a b. (a, b) -> a
fst) [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs =
      Loc -> Notes -> Doc () -> TermM (Pat ParamType)
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Loc
floc Notes
forall a. Monoid a => a
mempty (Doc () -> TermM (Pat ParamType))
-> Doc () -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Underscore-prefixed fields are not allowed."
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Did you mean"
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Text -> Doc ()
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Int -> Text -> Text
T.drop Int
1 (Name -> Text
nameToText Name
f)) Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"=_")
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"?"
  | [L Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([L Name] -> [L Name]
forall a. Ord a => [a] -> [a]
nubOrd (((L Name, PatBase NoInfo TyVar ParamType) -> L Name)
-> [(L Name, PatBase NoInfo TyVar ParamType)] -> [L Name]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, PatBase NoInfo TyVar ParamType) -> L Name
forall a b. (a, b) -> a
fst [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [L Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (((L Name, PatBase NoInfo TyVar ParamType) -> L Name)
-> [(L Name, PatBase NoInfo TyVar ParamType)] -> [L Name]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, PatBase NoInfo TyVar ParamType) -> L Name
forall a b. (a, b) -> a
fst [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs) =
      SrcLoc -> Notes -> Doc () -> TermM (Pat ParamType)
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM 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 () -> TermM (Pat ParamType))
-> Doc () -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Duplicate fields in record pattern" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> PatBase NoInfo TyVar ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. PatBase NoInfo TyVar ParamType -> Doc ann
pretty PatBase NoInfo TyVar ParamType
p Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
checkPat' p :: PatBase NoInfo TyVar ParamType
p@(RecordPat [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs SrcLoc
loc) (Ascribed ParamType
t)
  | Scalar (Record Map Name ParamType
t_fs) <- ParamType
t,
    [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs' <- ((L Name, PatBase NoInfo TyVar ParamType)
 -> (L Name, PatBase NoInfo TyVar ParamType) -> Ordering)
-> [(L Name, PatBase NoInfo TyVar ParamType)]
-> [(L Name, PatBase NoInfo TyVar ParamType)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy (((L Name, PatBase NoInfo TyVar ParamType) -> L Name)
-> (L Name, PatBase NoInfo TyVar ParamType)
-> (L Name, PatBase NoInfo TyVar ParamType)
-> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (L Name, PatBase NoInfo TyVar ParamType) -> L Name
forall a b. (a, b) -> a
fst) [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs,
    [(Name, ParamType)]
t_fs' <- ((Name, ParamType) -> (Name, ParamType) -> Ordering)
-> [(Name, ParamType)] -> [(Name, ParamType)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy (((Name, ParamType) -> Name)
-> (Name, ParamType) -> (Name, ParamType) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Name, ParamType) -> Name
forall a b. (a, b) -> a
fst) (Map Name ParamType -> [(Name, ParamType)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name ParamType
t_fs),
    ((Name, ParamType) -> Name) -> [(Name, ParamType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, ParamType) -> Name
forall a b. (a, b) -> a
fst [(Name, ParamType)]
t_fs' [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== ((L Name, PatBase NoInfo TyVar ParamType) -> Name)
-> [(L Name, PatBase NoInfo TyVar ParamType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (L Name -> Name
forall a. L a -> a
unLoc (L Name -> Name)
-> ((L Name, PatBase NoInfo TyVar ParamType) -> L Name)
-> (L Name, PatBase NoInfo TyVar ParamType)
-> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (L Name, PatBase NoInfo TyVar ParamType) -> L Name
forall a b. (a, b) -> a
fst) [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs' =
      [(L Name, Pat ParamType)] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([(L Name, Pat ParamType)] -> SrcLoc -> Pat ParamType)
-> TermM [(L Name, Pat ParamType)]
-> TermM (SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((L Name, PatBase NoInfo TyVar ParamType)
 -> (Name, ParamType) -> TermM (L Name, Pat ParamType))
-> [(L Name, PatBase NoInfo TyVar ParamType)]
-> [(Name, ParamType)]
-> TermM [(L Name, Pat ParamType)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (L Name, PatBase NoInfo TyVar ParamType)
-> (Name, ParamType) -> TermM (L Name, Pat ParamType)
forall {a} {a}.
(L a, PatBase NoInfo TyVar ParamType)
-> (a, ParamType) -> TermM (L a, Pat ParamType)
check [(L Name, PatBase NoInfo TyVar ParamType)]
p_fs' [(Name, ParamType)]
t_fs' TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  | Bool
otherwise = do
      SrcLoc -> Notes -> Doc () -> TermM (Pat ParamType)
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM 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 () -> TermM (Pat ParamType))
-> Doc () -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Pattern"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (PatBase NoInfo TyVar ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. PatBase NoInfo TyVar ParamType -> Doc ann
pretty PatBase NoInfo TyVar ParamType
p)
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"cannot match ascribed type"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ParamType -> Doc ann
pretty ParamType
t)
  where
    check :: (L a, PatBase NoInfo TyVar ParamType)
-> (a, ParamType) -> TermM (L a, Pat ParamType)
check (L Loc
f_loc a
f, PatBase NoInfo TyVar ParamType
p_f) (a
_, ParamType
t_f) =
      (Loc -> a -> L a
forall a. Loc -> a -> L a
L Loc
f_loc a
f,) (Pat ParamType -> (L a, Pat ParamType))
-> TermM (Pat ParamType) -> TermM (L a, Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' PatBase NoInfo TyVar ParamType
p_f (ParamType -> Inferred ParamType
forall t. t -> Inferred t
Ascribed ParamType
t_f)
checkPat' (RecordPat [(L Name, PatBase NoInfo TyVar ParamType)]
fs SrcLoc
loc) Inferred ParamType
NoneInferred =
  [(L Name, Pat ParamType)] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([(L Name, Pat ParamType)] -> SrcLoc -> Pat ParamType)
-> (Map (L Name) (Pat ParamType) -> [(L Name, Pat ParamType)])
-> Map (L Name) (Pat ParamType)
-> SrcLoc
-> Pat ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map (L Name) (Pat ParamType) -> [(L Name, Pat ParamType)]
forall k a. Map k a -> [(k, a)]
M.toList
    (Map (L Name) (Pat ParamType) -> SrcLoc -> Pat ParamType)
-> TermM (Map (L Name) (Pat ParamType))
-> TermM (SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase NoInfo TyVar ParamType -> TermM (Pat ParamType))
-> Map (L Name) (PatBase NoInfo TyVar ParamType)
-> TermM (Map (L Name) (Pat ParamType))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map (L Name) a -> f (Map (L Name) b)
traverse (PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
`checkPat'` Inferred ParamType
forall t. Inferred t
NoneInferred) ([(L Name, PatBase NoInfo TyVar ParamType)]
-> Map (L Name) (PatBase NoInfo TyVar ParamType)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(L Name, PatBase NoInfo TyVar ParamType)]
fs)
    TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkPat' (PatAscription PatBase NoInfo TyVar ParamType
p TypeExp (ExpBase NoInfo TyVar) TyVar
t SrcLoc
loc) Inferred ParamType
maybe_outer_t = do
  (t', _, RetType _ st, _) <- (ExpBase NoInfo TyVar -> TermM Exp)
-> TypeExp (ExpBase NoInfo TyVar) TyVar
-> TermM
     (TypeExp Exp TyVar, [TyVar], RetTypeBase Exp Freshness, Liftedness)
forall (m :: * -> *) df.
(MonadTypeChecker m, Pretty df) =>
(df -> m Exp)
-> TypeExp df TyVar
-> m (TypeExp Exp TyVar, [TyVar], RetTypeBase Exp Freshness,
      Liftedness)
checkTypeExp ExpBase NoInfo TyVar -> TermM Exp
checkSizeExp' TypeExp (ExpBase NoInfo TyVar) TyVar
t

  let st' = TypeBase Exp Freshness -> ParamType
resToParam TypeBase Exp Freshness
st

  case maybe_outer_t of
    Ascribed ParamType
outer_t -> do
      Bool -> TermM () -> TermM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ParamType -> TypeBase () Diet
forall o. TypeBase Exp o -> TypeBase () o
toType ParamType
st' TypeBase () Diet -> TypeBase () Diet -> Bool
forall a. Eq a => a -> a -> Bool
== ParamType -> TypeBase () Diet
forall o. TypeBase Exp o -> TypeBase () o
toType ParamType
outer_t) (TermM () -> TermM ()) -> TermM () -> TermM ()
forall a b. (a -> b) -> a -> b
$
        SrcLoc -> Notes -> Doc () -> TermM ()
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM 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 () -> TermM ()) -> Doc () -> TermM ()
forall a b. (a -> b) -> a -> b
$
          Doc ()
"Ascribed type"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (TypeBase Exp Freshness -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase Exp Freshness -> Doc ann
pretty TypeBase Exp Freshness
st)
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"cannot match outer ascribed type"
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ParamType -> Doc ann
pretty ParamType
outer_t)
      Pat ParamType -> TypeExp Exp TyVar -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
PatBase f vn t
-> TypeExp (ExpBase f vn) vn -> SrcLoc -> PatBase f vn t
PatAscription
        (Pat ParamType -> TypeExp Exp TyVar -> SrcLoc -> Pat ParamType)
-> TermM (Pat ParamType)
-> TermM (TypeExp Exp TyVar -> SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' PatBase NoInfo TyVar ParamType
p (ParamType -> Inferred ParamType
forall t. t -> Inferred t
Ascribed ParamType
st')
        TermM (TypeExp Exp TyVar -> SrcLoc -> Pat ParamType)
-> TermM (TypeExp Exp TyVar) -> TermM (SrcLoc -> Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeExp Exp TyVar -> TermM (TypeExp Exp TyVar)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp Exp TyVar
t'
        TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    Inferred ParamType
NoneInferred ->
      Pat ParamType -> TypeExp Exp TyVar -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
PatBase f vn t
-> TypeExp (ExpBase f vn) vn -> SrcLoc -> PatBase f vn t
PatAscription
        (Pat ParamType -> TypeExp Exp TyVar -> SrcLoc -> Pat ParamType)
-> TermM (Pat ParamType)
-> TermM (TypeExp Exp TyVar -> SrcLoc -> Pat ParamType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' PatBase NoInfo TyVar ParamType
p (ParamType -> Inferred ParamType
forall t. t -> Inferred t
Ascribed ParamType
st')
        TermM (TypeExp Exp TyVar -> SrcLoc -> Pat ParamType)
-> TermM (TypeExp Exp TyVar) -> TermM (SrcLoc -> Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeExp Exp TyVar -> TermM (TypeExp Exp TyVar)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp Exp TyVar
t'
        TermM (SrcLoc -> Pat ParamType)
-> TermM SrcLoc -> TermM (Pat ParamType)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkPat' (PatLit PatLit
l NoInfo ParamType
NoInfo SrcLoc
loc) (Ascribed ParamType
t) = do
  t' <- PatLit -> SrcLoc -> TermM ParamType
patLitMkType PatLit
l SrcLoc
loc
  ctEq (Reason (locOf loc)) (toType t') (toType t)
  pure $ PatLit l (Info t') loc
checkPat' (PatLit PatLit
l NoInfo ParamType
NoInfo SrcLoc
loc) Inferred ParamType
NoneInferred = do
  t' <- PatLit -> SrcLoc -> TermM ParamType
patLitMkType PatLit
l SrcLoc
loc
  pure $ PatLit l (Info t') loc
checkPat' (PatConstr Name
n NoInfo ParamType
NoInfo [PatBase NoInfo TyVar ParamType]
ps SrcLoc
loc) (Ascribed (Scalar (Sum Map Name [ParamType]
cs)))
  | Just [ParamType]
ts <- Name -> Map Name [ParamType] -> Maybe [ParamType]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n Map Name [ParamType]
cs,
    [PatBase NoInfo TyVar ParamType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PatBase NoInfo TyVar ParamType]
ps Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [ParamType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ParamType]
ts = do
      ps' <- (PatBase NoInfo TyVar ParamType
 -> Inferred ParamType -> TermM (Pat ParamType))
-> [PatBase NoInfo TyVar ParamType]
-> [Inferred ParamType]
-> TermM [Pat ParamType]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' [PatBase NoInfo TyVar ParamType]
ps ([Inferred ParamType] -> TermM [Pat ParamType])
-> [Inferred ParamType] -> TermM [Pat ParamType]
forall a b. (a -> b) -> a -> b
$ (ParamType -> Inferred ParamType)
-> [ParamType] -> [Inferred ParamType]
forall a b. (a -> b) -> [a] -> [b]
map ParamType -> Inferred ParamType
forall t. t -> Inferred t
Ascribed [ParamType]
ts
      pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc
checkPat' p :: PatBase NoInfo TyVar ParamType
p@(PatConstr {}) (Ascribed ParamType
t) =
  Loc -> Notes -> Doc () -> TermM (Pat ParamType)
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError (PatBase NoInfo TyVar ParamType -> Loc
forall a. Located a => a -> Loc
locOf PatBase NoInfo TyVar ParamType
p) Notes
forall a. Monoid a => a
mempty (Doc () -> TermM (Pat ParamType))
-> Doc () -> TermM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$
    Doc ()
"Pattern"
      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (PatBase NoInfo TyVar ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. PatBase NoInfo TyVar ParamType -> Doc ann
pretty PatBase NoInfo TyVar ParamType
p)
      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"cannot match ascribed type"
      Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ParamType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. ParamType -> Doc ann
pretty ParamType
t)
checkPat' (PatConstr Name
n NoInfo ParamType
NoInfo [PatBase NoInfo TyVar ParamType]
ps SrcLoc
loc) Inferred ParamType
NoneInferred = do
  ps' <- (PatBase NoInfo TyVar ParamType -> TermM (Pat ParamType))
-> [PatBase NoInfo TyVar ParamType] -> TermM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
`checkPat'` Inferred ParamType
forall t. Inferred t
NoneInferred) [PatBase NoInfo TyVar ParamType]
ps
  t <- newTypeWithConstr loc "t" Observe n =<< mapM (asType . patternType) ps'
  pure $ PatConstr n (Info $ toParam Observe t) ps' loc

checkPat ::
  PatBase NoInfo VName (TypeBase Size o) ->
  (Pat ParamType -> TermM a) ->
  TermM a
checkPat :: forall o a.
PatBase NoInfo TyVar (TypeBase Exp o)
-> (Pat ParamType -> TermM a) -> TermM a
checkPat PatBase NoInfo TyVar (TypeBase Exp o)
p Pat ParamType -> TermM a
m =
  Pat ParamType -> TermM a
m (Pat ParamType -> TermM a) -> TermM (Pat ParamType) -> TermM a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PatBase NoInfo TyVar ParamType
-> Inferred ParamType -> TermM (Pat ParamType)
checkPat' ((TypeBase Exp o -> ParamType)
-> PatBase NoInfo TyVar (TypeBase Exp o)
-> PatBase NoInfo TyVar ParamType
forall a b.
(a -> b) -> PatBase NoInfo TyVar a -> PatBase NoInfo TyVar b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Diet -> TypeBase Exp o -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe) PatBase NoInfo TyVar (TypeBase Exp o)
p) Inferred ParamType
forall t. Inferred t
NoneInferred

-- | Bind @let@-bound sizes. This is usually followed by 'bindLetPat'
-- immediately afterwards.
bindSizes :: [SizeBinder VName] -> TermM a -> TermM a
bindSizes :: forall a. [SizeBinder TyVar] -> TermM a -> TermM a
bindSizes [] TermM a
m = TermM a
m -- Minor optimisation.
bindSizes [SizeBinder TyVar]
sizes TermM a
m = [IdentBase Info TyVar StructType] -> TermM a -> TermM a
forall a. [IdentBase Info TyVar StructType] -> TermM a -> TermM a
bind ((SizeBinder TyVar -> IdentBase Info TyVar StructType)
-> [SizeBinder TyVar] -> [IdentBase Info TyVar StructType]
forall a b. (a -> b) -> [a] -> [b]
map SizeBinder TyVar -> IdentBase Info TyVar StructType
forall {vn} {dim} {o}.
SizeBinder vn -> IdentBase Info vn (TypeBase dim o)
sizeWithType [SizeBinder TyVar]
sizes) TermM a
m
  where
    sizeWithType :: SizeBinder vn -> IdentBase Info vn (TypeBase dim o)
sizeWithType SizeBinder vn
size =
      vn
-> Info (TypeBase dim o)
-> SrcLoc
-> IdentBase Info vn (TypeBase dim o)
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident (SizeBinder vn -> vn
forall vn. SizeBinder vn -> vn
sizeName SizeBinder vn
size) (TypeBase dim o -> Info (TypeBase dim o)
forall a. a -> Info a
Info (ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (IntType -> PrimType
Signed IntType
Int64)))) (SizeBinder vn -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf SizeBinder vn
size)

bindLetPat ::
  PatBase NoInfo VName (TypeBase Size o) ->
  Type ->
  (Pat ParamType -> TermM a) ->
  TermM a
bindLetPat :: forall o a.
PatBase NoInfo TyVar (TypeBase Exp o)
-> TypeBase () NoMode -> (Pat ParamType -> TermM a) -> TermM a
bindLetPat PatBase NoInfo TyVar (TypeBase Exp o)
p TypeBase () NoMode
t Pat ParamType -> TermM a
m = do
  PatBase NoInfo TyVar (TypeBase Exp o)
-> (Pat ParamType -> TermM a) -> TermM a
forall o a.
PatBase NoInfo TyVar (TypeBase Exp o)
-> (Pat ParamType -> TermM a) -> TermM a
checkPat PatBase NoInfo TyVar (TypeBase Exp o)
p ((Pat ParamType -> TermM a) -> TermM a)
-> (Pat ParamType -> TermM a) -> TermM a
forall a b. (a -> b) -> a -> b
$ \Pat ParamType
p' -> do
    pt <- ParamType -> TermM (TypeBase () Diet)
forall o. Monoid o => TypeBase Exp o -> TermM (TypeBase () o)
asType (ParamType -> TermM (TypeBase () Diet))
-> ParamType -> TermM (TypeBase () Diet)
forall a b. (a -> b) -> a -> b
$ Pat ParamType -> ParamType
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType Pat ParamType
p'
    ctEq (ReasonPatMatch (locOf p) (fmap toStruct p) t) pt t
    bind (patIdents (fmap toStruct p')) $ m p'

bindTypes ::
  [(VName, TypeBinding)] ->
  TermM a ->
  TermM a
bindTypes :: forall a. [(TyVar, TypeBinding)] -> TermM a -> TermM a
bindTypes [(TyVar, TypeBinding)]
tbinds = (TermScope () -> TermScope ()) -> TermM a -> TermM a
forall a. (TermScope () -> TermScope ()) -> TermM a -> TermM a
localScope TermScope () -> TermScope ()
forall {dim}. TermScope dim -> TermScope dim
extend
  where
    extend :: TermScope dim -> TermScope dim
extend TermScope dim
scope =
      TermScope dim
scope
        { scopeTypeTable = M.fromList tbinds <> scopeTypeTable scope
        }

bindTypeParams :: [TypeParam] -> TermM a -> TermM a
bindTypeParams :: forall a. [TypeParam] -> TermM a -> TermM a
bindTypeParams [TypeParam]
tparams TermM a
m =
  [IdentBase Info TyVar StructType] -> TermM a -> TermM a
forall a. [IdentBase Info TyVar StructType] -> TermM a -> TermM a
bind [IdentBase Info TyVar StructType]
idents (TermM a -> TermM a) -> (TermM a -> TermM a) -> TermM a -> TermM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TyVar, TypeBinding)] -> TermM a -> TermM a
forall a. [(TyVar, TypeBinding)] -> TermM a -> TermM a
bindTypes [(TyVar, TypeBinding)]
types (TermM a -> TermM a) -> TermM a -> TermM a
forall a b. (a -> b) -> a -> b
$ do
    lvl <- TermM Int
curLevel
    modify $ \TermState
s ->
      TermState
s
        { termTyParams =
            termTyParams s
              <> M.fromList (mapMaybe (typeParam lvl) tparams)
        }
    m
  where
    idents :: [IdentBase Info TyVar StructType]
idents = (TypeParam -> Maybe (IdentBase Info TyVar StructType))
-> [TypeParam] -> [IdentBase Info TyVar StructType]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe TypeParam -> Maybe (IdentBase Info TyVar StructType)
typeParamIdent [TypeParam]
tparams
    types :: [(TyVar, TypeBinding)]
types = (TypeParam -> Maybe (TyVar, TypeBinding))
-> [TypeParam] -> [(TyVar, TypeBinding)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe TypeParam -> Maybe (TyVar, TypeBinding)
typeParamType [TypeParam]
tparams
    typeParamType :: TypeParam -> Maybe (TyVar, TypeBinding)
typeParamType (TypeParamType Liftedness
l TyVar
v SrcLoc
_) =
      (TyVar, TypeBinding) -> Maybe (TyVar, TypeBinding)
forall a. a -> Maybe a
Just (TyVar
v, Liftedness -> [TypeParam] -> StructRetType -> TypeBinding
TypeAbbr Liftedness
l [] (StructRetType -> TypeBinding) -> StructRetType -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [TyVar] -> StructType -> StructRetType
forall dim o. [TyVar] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (StructType -> StructRetType) -> StructType -> StructRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode
-> QualName TyVar -> [TypeArg Exp] -> ScalarTypeBase Exp NoMode
forall dim o.
o -> QualName TyVar -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar NoMode
forall a. Monoid a => a
mempty (TyVar -> QualName TyVar
forall v. v -> QualName v
qualName TyVar
v) []))
    typeParamType TypeParamDim {} = Maybe (TyVar, TypeBinding)
forall a. Maybe a
Nothing
    typeParam :: a -> TypeParamBase a -> Maybe (a, (a, Liftedness, Loc))
typeParam a
lvl (TypeParamType Liftedness
l a
v SrcLoc
loc) = (a, (a, Liftedness, Loc)) -> Maybe (a, (a, Liftedness, Loc))
forall a. a -> Maybe a
Just (a
v, (a
lvl, Liftedness
l, SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc))
    typeParam a
_ TypeParamBase a
_ = Maybe (a, (a, Liftedness, Loc))
forall a. Maybe a
Nothing

bindParams ::
  [TypeParam] ->
  [PatBase NoInfo VName ParamType] ->
  ([Pat ParamType] -> TermM a) ->
  TermM a
bindParams :: forall a.
[TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType] -> TermM a)
-> TermM a
bindParams [TypeParam]
tps [PatBase NoInfo TyVar ParamType]
orig_ps [Pat ParamType] -> TermM a
m = [TypeParam] -> TermM a -> TermM a
forall a. [TypeParam] -> TermM a -> TermM a
bindTypeParams [TypeParam]
tps (TermM a -> TermM a) -> TermM a -> TermM a
forall a b. (a -> b) -> a -> b
$ do
  let descend :: [Pat ParamType]
-> [PatBase NoInfo TyVar (TypeBase Exp o)] -> TermM a
descend [Pat ParamType]
ps' (PatBase NoInfo TyVar (TypeBase Exp o)
p : [PatBase NoInfo TyVar (TypeBase Exp o)]
ps) =
        PatBase NoInfo TyVar (TypeBase Exp o)
-> (Pat ParamType -> TermM a) -> TermM a
forall o a.
PatBase NoInfo TyVar (TypeBase Exp o)
-> (Pat ParamType -> TermM a) -> TermM a
checkPat PatBase NoInfo TyVar (TypeBase Exp o)
p ((Pat ParamType -> TermM a) -> TermM a)
-> (Pat ParamType -> TermM a) -> TermM a
forall a b. (a -> b) -> a -> b
$ \Pat ParamType
p' ->
          [IdentBase Info TyVar StructType] -> TermM a -> TermM a
forall a. [IdentBase Info TyVar StructType] -> TermM a -> TermM a
bind (PatBase Info TyVar StructType -> [IdentBase Info TyVar StructType]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents (PatBase Info TyVar StructType
 -> [IdentBase Info TyVar StructType])
-> PatBase Info TyVar StructType
-> [IdentBase Info TyVar StructType]
forall a b. (a -> b) -> a -> b
$ (ParamType -> StructType)
-> Pat ParamType -> PatBase Info TyVar StructType
forall a b.
(a -> b) -> PatBase Info TyVar a -> PatBase Info TyVar b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct Pat ParamType
p') (TermM a -> TermM a) -> TermM a -> TermM a
forall a b. (a -> b) -> a -> b
$ TermM a -> TermM a
forall a. TermM a -> TermM a
incLevel (TermM a -> TermM a) -> TermM a -> TermM a
forall a b. (a -> b) -> a -> b
$ [Pat ParamType]
-> [PatBase NoInfo TyVar (TypeBase Exp o)] -> TermM a
descend (Pat ParamType
p' Pat ParamType -> [Pat ParamType] -> [Pat ParamType]
forall a. a -> [a] -> [a]
: [Pat ParamType]
ps') [PatBase NoInfo TyVar (TypeBase Exp o)]
ps
      descend [Pat ParamType]
ps' [] = [Pat ParamType] -> TermM a
m ([Pat ParamType] -> TermM a) -> [Pat ParamType] -> TermM a
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> [Pat ParamType]
forall a. [a] -> [a]
reverse [Pat ParamType]
ps'

  TermM a -> TermM a
forall a. TermM a -> TermM a
incLevel (TermM a -> TermM a) -> TermM a -> TermM a
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> [PatBase NoInfo TyVar ParamType] -> TermM a
forall {o}.
[Pat ParamType]
-> [PatBase NoInfo TyVar (TypeBase Exp o)] -> TermM a
descend [] [PatBase NoInfo TyVar ParamType]
orig_ps

checkApplyOne ::
  SrcLoc ->
  (Maybe (QualName VName), Int) ->
  Type ->
  (Maybe Exp, Type) ->
  TermM Type
checkApplyOne :: SrcLoc
-> (Maybe (QualName TyVar), Int)
-> TypeBase () NoMode
-> (Maybe Exp, TypeBase () NoMode)
-> TermM (TypeBase () NoMode)
checkApplyOne SrcLoc
loc (Maybe (QualName TyVar), Int)
fname TypeBase () NoMode
ftype (Maybe Exp
arg, TypeBase () NoMode
argtype) = do
  (a, b) <- TypeBase () NoMode
-> TermM (TypeBase () NoMode, TypeBase () NoMode)
split TypeBase () NoMode
ftype
  let reason = case Maybe Exp
arg of
        Just Exp
arg' -> Loc
-> (Maybe (QualName TyVar), Int)
-> Exp
-> TypeBase () NoMode
-> TypeBase () NoMode
-> Reason (TypeBase () NoMode)
forall t.
Loc -> (Maybe (QualName TyVar), Int) -> Exp -> t -> t -> Reason t
ReasonApply (Maybe Exp -> Loc
forall a. Located a => a -> Loc
locOf Maybe Exp
arg) (Maybe (QualName TyVar), Int)
fname Exp
arg' TypeBase () NoMode
a TypeBase () NoMode
argtype
        Maybe Exp
Nothing -> Loc -> Reason (TypeBase () NoMode)
forall t. Loc -> Reason t
Reason (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc)
  ctEq reason argtype a
  pure b
  where
    split :: TypeBase () NoMode
-> TermM (TypeBase () NoMode, TypeBase () NoMode)
split (Scalar (Arrow NoMode
_ PName
_ Diet
_ TypeBase () NoMode
a (RetType [TyVar]
_ TypeBase () Freshness
b))) =
      (TypeBase () NoMode, TypeBase () NoMode)
-> TermM (TypeBase () NoMode, TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase () NoMode
a, TypeBase () Freshness
b TypeBase () Freshness -> NoMode -> TypeBase () NoMode
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`setMode` NoMode
NoMode)
    split (Array NoMode
_u Shape ()
s ScalarTypeBase () NoMode
t) = do
      (a, b) <- TypeBase () NoMode
-> TermM (TypeBase () NoMode, TypeBase () NoMode)
split (TypeBase () NoMode
 -> TermM (TypeBase () NoMode, TypeBase () NoMode))
-> TypeBase () NoMode
-> TermM (TypeBase () NoMode, TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase () NoMode -> TypeBase () NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase () NoMode
t
      pure (arrayOf s a, arrayOf s b)
    split TypeBase () NoMode
ftype' = do
      a <- SrcLoc
-> Liftedness -> Name -> NoMode -> TermM (TypeBase () NoMode)
forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType SrcLoc
loc Liftedness
Lifted Name
"arg" NoMode
NoMode
      b <- newType loc Lifted "res" Nonfresh
      let reason = case Maybe Exp
arg of
            Just Exp
arg' -> Loc
-> (Maybe (QualName TyVar), Int)
-> Exp
-> TypeBase () NoMode
-> Reason (TypeBase () NoMode)
forall t.
Loc -> (Maybe (QualName TyVar), Int) -> Exp -> t -> Reason t
ReasonApplySplit (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Maybe (QualName TyVar), Int)
fname Exp
arg' TypeBase () NoMode
ftype'
            Maybe Exp
Nothing -> Loc -> Reason (TypeBase () NoMode)
forall t. Loc -> Reason t
Reason (Loc -> Reason (TypeBase () NoMode))
-> Loc -> Reason (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc
      ctEq reason ftype' $ Scalar $ Arrow NoMode Unnamed Observe a $ RetType [] b
      pure (a, b `setMode` NoMode)

checkApply ::
  SrcLoc ->
  Maybe (QualName VName) ->
  Type ->
  NE.NonEmpty (Maybe Exp, Type) ->
  TermM Type
checkApply :: SrcLoc
-> Maybe (QualName TyVar)
-> TypeBase () NoMode
-> NonEmpty (Maybe Exp, TypeBase () NoMode)
-> TermM (TypeBase () NoMode)
checkApply SrcLoc
loc Maybe (QualName TyVar)
fname TypeBase () NoMode
ftype NonEmpty (Maybe Exp, TypeBase () NoMode)
args = do
  (_, rt) <- ((Int, TypeBase () NoMode)
 -> (Maybe Exp, TypeBase () NoMode)
 -> TermM (Int, TypeBase () NoMode))
-> (Int, TypeBase () NoMode)
-> NonEmpty (Maybe Exp, TypeBase () NoMode)
-> TermM (Int, TypeBase () NoMode)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Int, TypeBase () NoMode)
-> (Maybe Exp, TypeBase () NoMode)
-> TermM (Int, TypeBase () NoMode)
onArg (Int
0, TypeBase () NoMode
ftype) NonEmpty (Maybe Exp, TypeBase () NoMode)
args
  pure rt
  where
    onArg :: (Int, TypeBase () NoMode)
-> (Maybe Exp, TypeBase () NoMode)
-> TermM (Int, TypeBase () NoMode)
onArg (Int
i, TypeBase () NoMode
f_t) (Maybe Exp, TypeBase () NoMode)
arg = do
      rt <- SrcLoc
-> (Maybe (QualName TyVar), Int)
-> TypeBase () NoMode
-> (Maybe Exp, TypeBase () NoMode)
-> TermM (TypeBase () NoMode)
checkApplyOne SrcLoc
loc (Maybe (QualName TyVar)
fname, Int
i) TypeBase () NoMode
f_t (Maybe Exp, TypeBase () NoMode)
arg
      pure (i + 1, rt)

checkSlice :: SliceBase NoInfo VName -> TermM [DimIndex]
checkSlice :: SliceBase NoInfo TyVar -> TermM [DimIndex]
checkSlice = (DimIndexBase NoInfo TyVar -> TermM DimIndex)
-> SliceBase NoInfo TyVar -> TermM [DimIndex]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM DimIndexBase NoInfo TyVar -> TermM DimIndex
checkDimIndex
  where
    checkDimIndex :: DimIndexBase NoInfo TyVar -> TermM DimIndex
checkDimIndex (DimFix ExpBase NoInfo TyVar
i) =
      Exp -> DimIndex
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (Exp -> DimIndex) -> TermM Exp -> TermM DimIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> [PrimType] -> Exp -> TermM Exp
require Text
"use as index" [PrimType]
anySignedType (Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
i)
    checkDimIndex (DimSlice Maybe (ExpBase NoInfo TyVar)
i Maybe (ExpBase NoInfo TyVar)
j Maybe (ExpBase NoInfo TyVar)
s) =
      Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndex
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndex)
-> TermM (Maybe Exp) -> TermM (Maybe Exp -> Maybe Exp -> DimIndex)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase NoInfo TyVar -> TermM Exp)
-> Maybe (ExpBase NoInfo TyVar) -> TermM (Maybe Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ExpBase NoInfo TyVar -> TermM Exp
check Maybe (ExpBase NoInfo TyVar)
i TermM (Maybe Exp -> Maybe Exp -> DimIndex)
-> TermM (Maybe Exp) -> TermM (Maybe Exp -> DimIndex)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (ExpBase NoInfo TyVar -> TermM Exp)
-> Maybe (ExpBase NoInfo TyVar) -> TermM (Maybe Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ExpBase NoInfo TyVar -> TermM Exp
check Maybe (ExpBase NoInfo TyVar)
j TermM (Maybe Exp -> DimIndex)
-> TermM (Maybe Exp) -> TermM DimIndex
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (ExpBase NoInfo TyVar -> TermM Exp)
-> Maybe (ExpBase NoInfo TyVar) -> TermM (Maybe Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ExpBase NoInfo TyVar -> TermM Exp
check Maybe (ExpBase NoInfo TyVar)
s

    check :: ExpBase NoInfo TyVar -> TermM Exp
check = Text -> [PrimType] -> Exp -> TermM Exp
require Text
"use in slice" [IntType -> PrimType
Signed IntType
Int64] (Exp -> TermM Exp)
-> (ExpBase NoInfo TyVar -> TermM Exp)
-> ExpBase NoInfo TyVar
-> TermM Exp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< ExpBase NoInfo TyVar -> TermM Exp
checkExp

isSlice :: DimIndexBase f vn -> Bool
isSlice :: forall (f :: * -> *) vn. DimIndexBase f vn -> Bool
isSlice DimSlice {} = Bool
True
isSlice DimFix {} = Bool
False

checkCase ::
  Type ->
  CaseBase NoInfo VName ->
  TermM (CaseBase Info VName, Type)
checkCase :: TypeBase () NoMode
-> CaseBase NoInfo TyVar
-> TermM (CaseBase Info TyVar, TypeBase () NoMode)
checkCase TypeBase () NoMode
mt (CasePat PatBase NoInfo TyVar StructType
p ExpBase NoInfo TyVar
e SrcLoc
loc) =
  PatBase NoInfo TyVar StructType
-> TypeBase () NoMode
-> (Pat ParamType
    -> TermM (CaseBase Info TyVar, TypeBase () NoMode))
-> TermM (CaseBase Info TyVar, TypeBase () NoMode)
forall o a.
PatBase NoInfo TyVar (TypeBase Exp o)
-> TypeBase () NoMode -> (Pat ParamType -> TermM a) -> TermM a
bindLetPat PatBase NoInfo TyVar StructType
p TypeBase () NoMode
mt ((Pat ParamType -> TermM (CaseBase Info TyVar, TypeBase () NoMode))
 -> TermM (CaseBase Info TyVar, TypeBase () NoMode))
-> (Pat ParamType
    -> TermM (CaseBase Info TyVar, TypeBase () NoMode))
-> TermM (CaseBase Info TyVar, TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ \Pat ParamType
p' -> do
    e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
    e_t <- expType e'
    pure (CasePat (fmap toStruct p') e' loc, e_t)

checkCases ::
  Type ->
  NE.NonEmpty (CaseBase NoInfo VName) ->
  TermM (NE.NonEmpty (CaseBase Info VName), Type)
checkCases :: TypeBase () NoMode
-> NonEmpty (CaseBase NoInfo TyVar)
-> TermM (NonEmpty (CaseBase Info TyVar), TypeBase () NoMode)
checkCases TypeBase () NoMode
mt NonEmpty (CaseBase NoInfo TyVar)
rest_cs = do
  let (CaseBase NoInfo TyVar
c, Maybe (NonEmpty (CaseBase NoInfo TyVar))
rest_cs') = NonEmpty (CaseBase NoInfo TyVar)
-> (CaseBase NoInfo TyVar,
    Maybe (NonEmpty (CaseBase NoInfo TyVar)))
forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (CaseBase NoInfo TyVar)
rest_cs
  (c', c_t) <- TypeBase () NoMode
-> CaseBase NoInfo TyVar
-> TermM (CaseBase Info TyVar, TypeBase () NoMode)
checkCase TypeBase () NoMode
mt CaseBase NoInfo TyVar
c
  case rest_cs' of
    Maybe (NonEmpty (CaseBase NoInfo TyVar))
Nothing ->
      (NonEmpty (CaseBase Info TyVar), TypeBase () NoMode)
-> TermM (NonEmpty (CaseBase Info TyVar), TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CaseBase Info TyVar -> NonEmpty (CaseBase Info TyVar)
forall a. a -> NonEmpty a
NE.singleton CaseBase Info TyVar
c', TypeBase () NoMode
c_t)
    Just NonEmpty (CaseBase NoInfo TyVar)
cs -> do
      (cs', cs_t) <- TypeBase () NoMode
-> NonEmpty (CaseBase NoInfo TyVar)
-> TermM (NonEmpty (CaseBase Info TyVar), TypeBase () NoMode)
checkCases TypeBase () NoMode
mt NonEmpty (CaseBase NoInfo TyVar)
cs
      ctEq (ReasonBranches (locOf c) c_t cs_t) c_t cs_t
      pure (NE.cons c' cs', c_t)

checkRetDecl ::
  Exp ->
  Maybe (TypeExp (ExpBase NoInfo VName) VName) ->
  TermM (Type, Maybe (TypeExp Exp VName))
checkRetDecl :: Exp
-> Maybe (TypeExp (ExpBase NoInfo TyVar) TyVar)
-> TermM (TypeBase () NoMode, Maybe (TypeExp Exp TyVar))
checkRetDecl Exp
body Maybe (TypeExp (ExpBase NoInfo TyVar) TyVar)
Nothing = (,Maybe (TypeExp Exp TyVar)
forall a. Maybe a
Nothing) (TypeBase () NoMode
 -> (TypeBase () NoMode, Maybe (TypeExp Exp TyVar)))
-> TermM (TypeBase () NoMode)
-> TermM (TypeBase () NoMode, Maybe (TypeExp Exp TyVar))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> TermM (TypeBase () NoMode)
expType Exp
body
checkRetDecl Exp
body (Just TypeExp (ExpBase NoInfo TyVar) TyVar
te) = do
  (te', _, RetType _ st, _) <- (ExpBase NoInfo TyVar -> TermM Exp)
-> TypeExp (ExpBase NoInfo TyVar) TyVar
-> TermM
     (TypeExp Exp TyVar, [TyVar], RetTypeBase Exp Freshness, Liftedness)
forall (m :: * -> *) df.
(MonadTypeChecker m, Pretty df) =>
(df -> m Exp)
-> TypeExp df TyVar
-> m (TypeExp Exp TyVar, [TyVar], RetTypeBase Exp Freshness,
      Liftedness)
checkTypeExp ExpBase NoInfo TyVar -> TermM Exp
checkSizeExp' TypeExp (ExpBase NoInfo TyVar) TyVar
te
  body_t <- expType body
  st' <- toStruct <$> asType st
  ctEq (ReasonRetType (locOf body) st' body_t) st' body_t
  pure (st', Just te')

-- Add constraints saying that the first type has a (potentially nested) part
-- containing the second type.
--
-- FIXME: the locations here are very bad.
mustHaveSteps ::
  (Pretty a, Located a) =>
  a ->
  Type ->
  [UpdateStep Info VName] ->
  Type ->
  TermM ()
mustHaveSteps :: forall a.
(Pretty a, Located a) =>
a
-> TypeBase () NoMode
-> [UpdateStep Info TyVar]
-> TypeBase () NoMode
-> TermM ()
mustHaveSteps a
src TypeBase () NoMode
t [] TypeBase () NoMode
ve_t =
  -- This case is probably never reached.
  Reason (TypeBase () NoMode)
-> TypeBase () NoMode -> TypeBase () NoMode -> TermM ()
forall o1 o2.
Reason (TypeBase () NoMode)
-> TypeBase () o1 -> TypeBase () o2 -> TermM ()
ctEq (Loc -> Reason (TypeBase () NoMode)
forall t. Loc -> Reason t
Reason (a -> Loc
forall a. Located a => a -> Loc
locOf a
src)) TypeBase () NoMode
t TypeBase () NoMode
ve_t
mustHaveSteps a
src TypeBase () NoMode
t [UpdateStepField Name
f] TypeBase () NoMode
ve_t = do
  rt :: Type <- SrcLoc
-> Name -> Name -> TypeBase () NoMode -> TermM (TypeBase () NoMode)
newTypeWithField (a -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf a
src) Name
"ft" Name
f TypeBase () NoMode
ve_t
  ctEq (Reason (locOf src)) t rt
mustHaveSteps a
src TypeBase () NoMode
t (UpdateStepField Name
f : [UpdateStep Info TyVar]
steps) TypeBase () NoMode
ve_t = do
  ft <- Loc -> Liftedness -> Name -> NoMode -> TermM (TypeBase () NoMode)
forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType (a -> Loc
forall a. Located a => a -> Loc
locOf a
src) Liftedness
Lifted Name
"ft" NoMode
NoMode
  rt :: Type <- newTypeWithField (srclocOf src) "ft" f ft
  ctEq (Reason (locOf src)) t rt
  mustHaveSteps src ft steps ve_t
mustHaveSteps a
src TypeBase () NoMode
t [UpdateStepSlice [DimIndex]
slice] TypeBase () NoMode
ve_t = do
  let num_slices :: Int
num_slices = [DimIndex] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DimIndex] -> Int) -> [DimIndex] -> Int
forall a b. (a -> b) -> a -> b
$ (DimIndex -> Bool) -> [DimIndex] -> [DimIndex]
forall a. (a -> Bool) -> [a] -> [a]
filter DimIndex -> Bool
forall (f :: * -> *) vn. DimIndexBase f vn -> Bool
isSlice [DimIndex]
slice
  update_elem_t <- Loc -> Name -> NoMode -> TermM (TypeBase () NoMode)
forall loc o dim.
Located loc =>
loc -> Name -> o -> TermM (TypeBase dim o)
newElemType (a -> Loc
forall a. Located a => a -> Loc
locOf a
src) Name
"update_elem" NoMode
NoMode
  ctEq (Reason (locOf src)) t $ arrayOfRank (length slice) update_elem_t
  ctEq (Reason (locOf src)) ve_t $ arrayOfRank num_slices update_elem_t
mustHaveSteps a
src TypeBase () NoMode
t (UpdateStepSlice [DimIndex]
slice : [UpdateStep Info TyVar]
steps) TypeBase () NoMode
ve_t = do
  let num_slices :: Int
num_slices = [DimIndex] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DimIndex] -> Int) -> [DimIndex] -> Int
forall a b. (a -> b) -> a -> b
$ (DimIndex -> Bool) -> [DimIndex] -> [DimIndex]
forall a. (a -> Bool) -> [a] -> [a]
filter DimIndex -> Bool
forall (f :: * -> *) vn. DimIndexBase f vn -> Bool
isSlice [DimIndex]
slice
  index_tv <- Loc -> Liftedness -> Name -> TermM TyVar
forall loc. Located loc => loc -> Liftedness -> Name -> TermM TyVar
newTyVar (a -> Loc
forall a. Located a => a -> Loc
locOf a
src) Liftedness
Unlifted Name
"index"
  index_elem_t <- newElemType (locOf src) "index_elem" NoMode
  ctEq (Reason (locOf src)) (tyVarType NoMode index_tv) $ arrayOfRank num_slices index_elem_t
  ctEq (Reason (locOf src)) t $ arrayOfRank (length slice) index_elem_t
  mustHaveSteps src (arrayOfRank num_slices index_elem_t) steps ve_t

checkStep :: UpdateStep NoInfo VName -> TermM (UpdateStep Info VName)
checkStep :: UpdateStep NoInfo TyVar -> TermM (UpdateStep Info TyVar)
checkStep (UpdateStepField Name
f) = UpdateStep Info TyVar -> TermM (UpdateStep Info TyVar)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep Info TyVar -> TermM (UpdateStep Info TyVar))
-> UpdateStep Info TyVar -> TermM (UpdateStep Info TyVar)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep Info TyVar
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
checkStep (UpdateStepSlice SliceBase NoInfo TyVar
slice) = [DimIndex] -> UpdateStep Info TyVar
forall (f :: * -> *) vn. SliceBase f vn -> UpdateStep f vn
UpdateStepSlice ([DimIndex] -> UpdateStep Info TyVar)
-> TermM [DimIndex] -> TermM (UpdateStep Info TyVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SliceBase NoInfo TyVar -> TermM [DimIndex]
checkSlice SliceBase NoInfo TyVar
slice

checkExp :: ExpBase NoInfo VName -> TermM (ExpBase Info VName)
--
checkExp :: ExpBase NoInfo TyVar -> TermM Exp
checkExp (Var QualName TyVar
qn NoInfo StructType
_ SrcLoc
loc) = do
  t <- TypeBase () NoMode -> TermM StructType
forall o. TypeBase () o -> TermM (TypeBase Exp o)
asStructType (TypeBase () NoMode -> TermM StructType)
-> TermM (TypeBase () NoMode) -> TermM StructType
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
loc QualName TyVar
qn
  pure $ Var qn (Info t) loc
checkExp (OpSection QualName TyVar
op NoInfo StructType
_ SrcLoc
loc) = do
  ftype <- TypeBase () NoMode -> TermM StructType
forall o. TypeBase () o -> TermM (TypeBase Exp o)
asStructType (TypeBase () NoMode -> TermM StructType)
-> TermM (TypeBase () NoMode) -> TermM StructType
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
loc QualName TyVar
op
  pure $ OpSection op (Info ftype) loc
checkExp (Negate ExpBase NoInfo TyVar
arg SrcLoc
loc) = do
  arg' <- Text -> [PrimType] -> Exp -> TermM Exp
require Text
"numeric negation" [PrimType]
anyNumberType (Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
arg
  pure $ Negate arg' loc
checkExp (Not ExpBase NoInfo TyVar
arg SrcLoc
loc) = do
  arg' <- Text -> [PrimType] -> Exp -> TermM Exp
require Text
"logical negation" (PrimType
Bool PrimType -> [PrimType] -> [PrimType]
forall a. a -> [a] -> [a]
: [PrimType]
anyIntType) (Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
arg
  pure $ Not arg' loc
checkExp (Hole NoInfo StructType
NoInfo SrcLoc
loc) =
  Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (Info StructType -> SrcLoc -> Exp)
-> TermM (Info StructType) -> TermM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType)
-> TermM StructType -> TermM (Info StructType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc -> Liftedness -> Name -> NoMode -> TermM StructType
forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType SrcLoc
loc Liftedness
Lifted Name
"hole" NoMode
NoMode) TermM (SrcLoc -> Exp) -> TermM SrcLoc -> TermM Exp
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkExp (Parens ExpBase NoInfo TyVar
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp) -> TermM Exp -> TermM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e TermM (SrcLoc -> Exp) -> TermM SrcLoc -> TermM Exp
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkExp (TupLit [ExpBase NoInfo TyVar]
es SrcLoc
loc) =
  [Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Exp] -> SrcLoc -> Exp) -> TermM [Exp] -> TermM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase NoInfo TyVar -> TermM Exp)
-> [ExpBase NoInfo TyVar] -> TermM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ExpBase NoInfo TyVar -> TermM Exp
checkExp [ExpBase NoInfo TyVar]
es TermM (SrcLoc -> Exp) -> TermM SrcLoc -> TermM Exp
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkExp (QualParens (QualName TyVar
modname, SrcLoc
modnameloc) ExpBase NoInfo TyVar
e SrcLoc
loc) = do
  mod <- QualName TyVar -> TermM Mod
lookupMod QualName TyVar
modname
  case mod of
    ModEnv Env
env -> (TermEnv -> TermEnv) -> TermM Exp -> TermM Exp
forall a. (TermEnv -> TermEnv) -> TermM a -> TermM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (TermEnv -> Env -> TermEnv
`withEnv` Env
env) (TermM Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall a b. (a -> b) -> a -> b
$ do
      e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
      pure $ QualParens (modname, modnameloc) e' loc
    ModFun {} ->
      SrcLoc -> Notes -> Doc () -> TermM Exp
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM 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 () -> TermM Exp) -> (Doc () -> Doc ()) -> Doc () -> TermM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"module-is-parametric" (Doc () -> TermM Exp) -> Doc () -> TermM Exp
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Module" Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> QualName TyVar -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName TyVar -> Doc ann
pretty QualName TyVar
modname Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
" is a parametric module."
--
checkExp (IntLit Integer
x NoInfo StructType
NoInfo SrcLoc
loc) = do
  t <- SrcLoc -> Name -> [PrimType] -> TermM StructType
forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded SrcLoc
loc Name
"num" [PrimType]
anyNumberType
  pure $ IntLit x (Info t) loc
checkExp (FloatLit Double
x NoInfo StructType
NoInfo SrcLoc
loc) = do
  t <- SrcLoc -> Name -> [PrimType] -> TermM StructType
forall d. SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoMode)
newTypeOverloaded SrcLoc
loc Name
"float" [PrimType]
anyFloatType
  pure $ FloatLit x (Info t) loc
checkExp (Literal PrimValue
v SrcLoc
loc) =
  Exp -> TermM Exp
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermM Exp) -> Exp -> TermM Exp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal PrimValue
v SrcLoc
loc
checkExp (StringLit [Word8]
vs SrcLoc
loc) =
  Exp -> TermM Exp
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermM Exp) -> Exp -> TermM Exp
forall a b. (a -> b) -> a -> b
$ [Word8] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [Word8] -> SrcLoc -> ExpBase f vn
StringLit [Word8]
vs SrcLoc
loc
-- No need to type check this, as these are only produced by the
-- parser if the elements are monomorphic and all match.
checkExp (ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  Exp -> TermM Exp
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> TermM Exp) -> Exp -> TermM Exp
forall a b. (a -> b) -> a -> b
$ [PrimValue] -> PrimType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc
checkExp (ArrayLit [ExpBase NoInfo TyVar]
es NoInfo StructType
_ SrcLoc
loc) = do
  -- TODO: this will produce an enormous number of constraints and
  -- type variables for pathologically large arrays with
  -- type-unsuffixed integers. Add some special case that handles that
  -- more efficiently.
  et <- SrcLoc -> Name -> NoMode -> TermM StructType
forall loc o dim.
Located loc =>
loc -> Name -> o -> TermM (TypeBase dim o)
newElemType SrcLoc
loc Name
"et" NoMode
NoMode
  es' <- forM es $ \ExpBase NoInfo TyVar
e -> do
    e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
    e_t <- expType e'
    et' <- asType et
    ctEq (Reason (locOf loc)) e_t et'
    pure e'
  let arr_t = Shape Exp -> StructType -> StructType
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Exp
sizeFromInteger ([ExpBase NoInfo TyVar] -> Integer
forall i a. Num i => [a] -> i
L.genericLength [ExpBase NoInfo TyVar]
es) SrcLoc
loc]) StructType
et
  pure $ ArrayLit es' (Info arr_t) loc
checkExp (RecordLit [FieldBase NoInfo TyVar]
fs SrcLoc
loc) =
  [FieldBase Info TyVar] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info TyVar] -> SrcLoc -> Exp)
-> TermM [FieldBase Info TyVar] -> TermM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (Map Name SrcLoc) TermM [FieldBase Info TyVar]
-> Map Name SrcLoc -> TermM [FieldBase Info TyVar]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT ((FieldBase NoInfo TyVar
 -> StateT (Map Name SrcLoc) TermM (FieldBase Info TyVar))
-> [FieldBase NoInfo TyVar]
-> StateT (Map Name SrcLoc) TermM [FieldBase Info TyVar]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM FieldBase NoInfo TyVar
-> StateT (Map Name SrcLoc) TermM (FieldBase Info TyVar)
forall {t :: (* -> *) -> * -> *}.
(MonadState (Map Name SrcLoc) (t TermM), MonadTrans t) =>
FieldBase NoInfo TyVar -> t TermM (FieldBase Info TyVar)
checkField [FieldBase NoInfo TyVar]
fs) Map Name SrcLoc
forall a. Monoid a => a
mempty TermM (SrcLoc -> Exp) -> TermM SrcLoc -> TermM Exp
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  where
    checkField :: FieldBase NoInfo TyVar -> t TermM (FieldBase Info TyVar)
checkField (RecordFieldExplicit L Name
f ExpBase NoInfo TyVar
e SrcLoc
rloc) = do
      Name -> SrcLoc -> t TermM ()
forall {a} {b} {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadState (Map a b) (t m), Ord a, MonadTrans t,
 MonadTypeChecker m, Pretty a, Located a, Located b) =>
a -> a -> t m ()
errIfAlreadySet (L Name -> Name
forall a. L a -> a
unLoc L Name
f) SrcLoc
rloc
      (Map Name SrcLoc -> Map Name SrcLoc) -> t TermM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map Name SrcLoc -> Map Name SrcLoc) -> t TermM ())
-> (Map Name SrcLoc -> Map Name SrcLoc) -> t TermM ()
forall a b. (a -> b) -> a -> b
$ Name -> SrcLoc -> Map Name SrcLoc -> Map Name SrcLoc
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (L Name -> Name
forall a. L a -> a
unLoc L Name
f) SrcLoc
rloc
      L Name -> Exp -> SrcLoc -> FieldBase Info TyVar
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit L Name
f (Exp -> SrcLoc -> FieldBase Info TyVar)
-> t TermM Exp -> t TermM (SrcLoc -> FieldBase Info TyVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TermM Exp -> t TermM Exp
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e) t TermM (SrcLoc -> FieldBase Info TyVar)
-> t TermM SrcLoc -> t TermM (FieldBase Info TyVar)
forall a b. t TermM (a -> b) -> t TermM a -> t TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t TermM SrcLoc
forall a. a -> t TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
rloc
    checkField (RecordFieldImplicit L TyVar
name NoInfo StructType
NoInfo SrcLoc
rloc) = do
      Name -> SrcLoc -> t TermM ()
forall {a} {b} {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadState (Map a b) (t m), Ord a, MonadTrans t,
 MonadTypeChecker m, Pretty a, Located a, Located b) =>
a -> a -> t m ()
errIfAlreadySet (TyVar -> Name
baseName (L TyVar -> TyVar
forall a. L a -> a
unLoc L TyVar
name)) SrcLoc
rloc
      t <- TermM StructType -> t TermM StructType
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (TermM StructType -> t TermM StructType)
-> TermM StructType -> t TermM StructType
forall a b. (a -> b) -> a -> b
$ TypeBase () NoMode -> TermM StructType
forall o. TypeBase () o -> TermM (TypeBase Exp o)
asStructType (TypeBase () NoMode -> TermM StructType)
-> TermM (TypeBase () NoMode) -> TermM StructType
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
rloc (TyVar -> QualName TyVar
forall v. v -> QualName v
qualName (L TyVar -> TyVar
forall a. L a -> a
unLoc L TyVar
name))
      modify $ M.insert (baseName (unLoc name)) rloc
      pure $ RecordFieldImplicit name (Info t) rloc

    errIfAlreadySet :: a -> a -> t m ()
errIfAlreadySet a
f a
rloc = do
      maybe_sloc <- (Map a b -> Maybe b) -> t m (Maybe b)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Map a b -> Maybe b) -> t m (Maybe b))
-> (Map a b -> Maybe b) -> t m (Maybe b)
forall a b. (a -> b) -> a -> b
$ a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
f
      case maybe_sloc of
        Just b
sloc ->
          m () -> t m ()
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> t m ()) -> (Doc () -> m ()) -> Doc () -> t m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Notes -> Doc () -> m ()
forall loc a. Located loc => loc -> Notes -> Doc () -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError a
rloc Notes
forall a. Monoid a => a
mempty (Doc () -> t m ()) -> Doc () -> t m ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Field"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (a -> Doc ()
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
f)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> Doc ()
"previously defined at"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> [Char] -> Doc ()
forall ann. [Char] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (a -> b -> [Char]
forall a b. (Located a, Located b) => a -> b -> [Char]
locStrRel a
rloc b
sloc)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
        Maybe b
Nothing -> () -> t m ()
forall a. a -> t m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

--
checkExp (Attr AttrInfo TyVar
info ExpBase NoInfo TyVar
e SrcLoc
loc) =
  AttrInfo TyVar -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr (AttrInfo TyVar -> Exp -> SrcLoc -> Exp)
-> TermM (AttrInfo TyVar) -> TermM (Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AttrInfo TyVar -> TermM (AttrInfo TyVar)
forall (m :: * -> *).
MonadTypeChecker m =>
AttrInfo TyVar -> m (AttrInfo TyVar)
checkAttr AttrInfo TyVar
info TermM (Exp -> SrcLoc -> Exp) -> TermM Exp -> TermM (SrcLoc -> Exp)
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e TermM (SrcLoc -> Exp) -> TermM SrcLoc -> TermM Exp
forall a b. TermM (a -> b) -> TermM a -> TermM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TermM SrcLoc
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
checkExp (Assert ExpBase NoInfo TyVar
e1 ExpBase NoInfo TyVar
e2 NoInfo Text
NoInfo SrcLoc
loc) = do
  e1' <- Text -> [PrimType] -> Exp -> TermM Exp
require Text
"being asserted" [PrimType
Bool] (Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e1
  e2' <- checkExp e2
  pure $ Assert e1' e2' (Info (prettyText e1)) loc
--
checkExp (Constr Name
name [ExpBase NoInfo TyVar]
es NoInfo StructType
NoInfo SrcLoc
loc) = do
  es' <- (ExpBase NoInfo TyVar -> TermM Exp)
-> [ExpBase NoInfo TyVar] -> TermM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ExpBase NoInfo TyVar -> TermM Exp
checkExp [ExpBase NoInfo TyVar]
es
  es_ts <- mapM expType es'
  t <- newTypeWithConstr loc "t" NoMode name es_ts
  pure $ Constr name es' (Info t) loc
--
checkExp (AppExp (Apply ExpBase NoInfo TyVar
fe NonEmpty (NoInfo (Maybe TyVar), ExpBase NoInfo TyVar)
args SrcLoc
loc) NoInfo AppRes
NoInfo) = do
  fe' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
fe
  (args', apply_args) <-
    fmap NE.unzip . forM args $ \(NoInfo (Maybe TyVar)
_, ExpBase NoInfo TyVar
arg) -> do
      arg' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
arg
      arg_t <- expType arg'
      pure (arg', (Just arg', arg_t))
  fe_t <- expType fe'
  rt <- checkApply loc fname fe_t apply_args
  rt' <- asStructType rt
  let args'' = (Exp -> (Info (Maybe a), Exp))
-> NonEmpty Exp -> NonEmpty (Info (Maybe a), Exp)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NE.map (\Exp
arg -> (Maybe a -> Info (Maybe a)
forall a. a -> Info a
Info Maybe a
forall a. Maybe a
Nothing, Exp
arg)) NonEmpty Exp
args'
  pure $ AppExp (Apply fe' args'' loc) $ Info (AppRes rt' [])
  where
    fname :: Maybe (QualName TyVar)
fname =
      case ExpBase NoInfo TyVar
fe of
        Var QualName TyVar
v NoInfo StructType
_ SrcLoc
_ -> QualName TyVar -> Maybe (QualName TyVar)
forall a. a -> Maybe a
Just QualName TyVar
v
        ExpBase NoInfo TyVar
_ -> Maybe (QualName TyVar)
forall a. Maybe a
Nothing
checkExp (AppExp (BinOp (QualName TyVar
op, SrcLoc
oploc) NoInfo StructType
NoInfo (ExpBase NoInfo TyVar
e1, NoInfo (Maybe TyVar)
_) (ExpBase NoInfo TyVar
e2, NoInfo (Maybe TyVar)
_) SrcLoc
loc) NoInfo AppRes
NoInfo) = do
  ftype <- SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
oploc QualName TyVar
op
  e1' <- checkExp e1
  e1_t <- expType e1'
  e2' <- checkExp e2
  e2_t <- expType e2'

  rt <-
    checkApply
      loc
      (Just op)
      ftype
      ((Just e1', e1_t) NE.:| [(Just e2', e2_t)])
  rt' <- asStructType rt

  ftype' <- asStructType ftype
  pure $
    AppExp
      (BinOp (op, oploc) (Info ftype') (e1', Info Nothing) (e2', Info Nothing) loc)
      (Info (AppRes rt' []))
--
checkExp (OpSectionLeft QualName TyVar
op NoInfo StructType
_ ExpBase NoInfo TyVar
e (NoInfo (PName, ParamType, Maybe TyVar), NoInfo (PName, ParamType))
_ (NoInfo (RetTypeBase Exp Freshness), NoInfo [TyVar])
_ SrcLoc
loc) = do
  optype <- SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
loc QualName TyVar
op
  e' <- checkExp e
  e_t <- expType e'
  t2 <- newType loc Lifted "t" NoMode
  t2' <- asStructType t2
  rt <-
    checkApply
      loc
      (Just op)
      optype
      ((Just e', e_t) NE.:| [(Nothing, t2)])
  rt' <- asStructType rt

  t1 <- asStructType e_t
  optype' <- asStructType optype
  pure $
    OpSectionLeft
      op
      (Info optype')
      e'
      ( Info (Unnamed, toParam Observe t1, Nothing),
        Info (Unnamed, toParam Observe t2')
      )
      (Info (RetType [] (rt' `setMode` Nonfresh)), Info [])
      loc
checkExp (OpSectionRight QualName TyVar
op NoInfo StructType
_ ExpBase NoInfo TyVar
e (NoInfo (PName, ParamType), NoInfo (PName, ParamType, Maybe TyVar))
_ NoInfo (RetTypeBase Exp Freshness)
NoInfo SrcLoc
loc) = do
  optype <- SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar SrcLoc
loc QualName TyVar
op
  e' <- checkExp e
  e_t <- expType e'
  t1 <- newType loc Lifted "t" NoMode
  t1' <- asStructType t1
  rt <-
    checkApply
      loc
      (Just op)
      optype
      ((Nothing, t1) NE.:| [(Just e', e_t)])
  rt' <- asStructType rt
  t2 <- asStructType e_t

  optype' <- asStructType optype
  pure $
    OpSectionRight
      op
      (Info optype')
      e'
      -- Dummy types.
      ( Info (Unnamed, toParam Observe t1'),
        Info (Unnamed, toParam Observe t2, Nothing)
      )
      (Info $ RetType [] (rt' `setMode` Nonfresh))
      loc
--
checkExp e :: ExpBase NoInfo TyVar
e@(UpdateSection [UpdateStep NoInfo TyVar]
steps NoInfo StructType
NoInfo SrcLoc
loc) = do
  steps' <- (UpdateStep NoInfo TyVar -> TermM (UpdateStep Info TyVar))
-> [UpdateStep NoInfo TyVar] -> TermM [UpdateStep Info TyVar]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM UpdateStep NoInfo TyVar -> TermM (UpdateStep Info TyVar)
checkStep [UpdateStep NoInfo TyVar]
steps
  -- Lifted, as a pure field projection works on records with
  -- function-typed fields. Any slice steps will constrain the
  -- relevant parts to be arrays (of unlifted elements) anyway.
  src_t <- newType loc Lifted "update" NoMode
  ve_t <- newType loc Lifted "update_elem" NoMode
  mustHaveSteps e src_t steps' ve_t
  ft <-
    asStructType $
      Scalar $
        Arrow mempty Unnamed Observe src_t $
          second (const Nonfresh) (RetType [] ve_t)
  pure $ UpdateSection steps' (Info ft) loc

--
checkExp (Lambda [PatBase NoInfo TyVar ParamType]
params ExpBase NoInfo TyVar
body Maybe (TypeExp (ExpBase NoInfo TyVar) TyVar)
retdecl NoInfo (RetTypeBase Exp Freshness)
NoInfo SrcLoc
loc) = do
  [TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType] -> TermM Exp)
-> TermM Exp
forall a.
[TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType] -> TermM a)
-> TermM a
bindParams [] [PatBase NoInfo TyVar ParamType]
params (([Pat ParamType] -> TermM Exp) -> TermM Exp)
-> ([Pat ParamType] -> TermM Exp) -> TermM Exp
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
    body' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
body

    (body_t, retdecl') <- checkRetDecl body' retdecl
    body_t' <- asStructType body_t
    let ret = [TyVar] -> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall dim o. [TyVar] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (TypeBase Exp Freshness -> RetTypeBase Exp Freshness)
-> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> TypeBase Exp Freshness
forall o. Freshness -> TypeBase Exp o -> TypeBase Exp Freshness
toRes Freshness
Nonfresh StructType
body_t'
    pure $ Lambda params' body' retdecl' (Info ret) loc
--
checkExp (AppExp (LetPat [SizeBinder TyVar]
sizes PatBase NoInfo TyVar StructType
pat ExpBase NoInfo TyVar
e ExpBase NoInfo TyVar
body SrcLoc
loc) NoInfo AppRes
_) = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  e_t <- expType e'

  bindSizes sizes . incLevel . bindLetPat pat e_t $ \Pat ParamType
pat' -> do
    body' <- TermM Exp -> TermM Exp
forall a. TermM a -> TermM a
incLevel (TermM Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall a b. (a -> b) -> a -> b
$ ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
body
    body_t <- expType body'

    body_t' <- asStructType body_t
    pure $
      AppExp
        (LetPat sizes (fmap toStruct pat') e' body' loc)
        (Info $ AppRes body_t' [])
--
checkExp (AppExp (LetFun (TyVar, SrcLoc)
name ([TypeParam]
tparams, [PatBase NoInfo TyVar ParamType]
params, Maybe (TypeExp (ExpBase NoInfo TyVar) TyVar)
retdecl, NoInfo (RetTypeBase Exp Freshness)
NoInfo, ExpBase NoInfo TyVar
e) ExpBase NoInfo TyVar
body SrcLoc
loc) NoInfo AppRes
_) = do
  (tparams', params', retdecl', rettype, e') <-
    [TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType]
    -> TermM
         ([TypeParam], [Pat ParamType], Maybe (TypeExp Exp TyVar),
          TypeBase () Freshness, Exp))
-> TermM
     ([TypeParam], [Pat ParamType], Maybe (TypeExp Exp TyVar),
      TypeBase () Freshness, Exp)
forall a.
[TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType] -> TermM a)
-> TermM a
bindParams [TypeParam]
tparams [PatBase NoInfo TyVar ParamType]
params (([Pat ParamType]
  -> TermM
       ([TypeParam], [Pat ParamType], Maybe (TypeExp Exp TyVar),
        TypeBase () Freshness, Exp))
 -> TermM
      ([TypeParam], [Pat ParamType], Maybe (TypeExp Exp TyVar),
       TypeBase () Freshness, Exp))
-> ([Pat ParamType]
    -> TermM
         ([TypeParam], [Pat ParamType], Maybe (TypeExp Exp TyVar),
          TypeBase () Freshness, Exp))
-> TermM
     ([TypeParam], [Pat ParamType], Maybe (TypeExp Exp TyVar),
      TypeBase () Freshness, Exp)
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
      e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
      (e_t, retdecl') <- checkRetDecl e' retdecl
      pure (tparams, params', retdecl', fmap (const Nonfresh) e_t, e')

  params'' <- mapM (traverse asType) params'

  let entry = [TypeParam] -> TypeBase () NoMode -> ValBinding ()
forall dim. [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
BoundV [TypeParam]
tparams' (TypeBase () NoMode -> ValBinding ())
-> TypeBase () NoMode -> ValBinding ()
forall a b. (a -> b) -> a -> b
$ [PatBase Info TyVar (TypeBase () Diet)]
-> RetTypeBase () Freshness -> TypeBase () NoMode
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [PatBase Info TyVar (TypeBase () Diet)]
params'' (RetTypeBase () Freshness -> TypeBase () NoMode)
-> RetTypeBase () Freshness -> TypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ [TyVar] -> TypeBase () Freshness -> RetTypeBase () Freshness
forall dim o. [TyVar] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase () Freshness
rettype
      bindF TermScope ()
scope =
        TermScope ()
scope
          { scopeVtable = M.insert (fst name) entry $ scopeVtable scope
          }
  body' <- localScope bindF $ checkExp body
  body_t <- expType body'

  body_t' <- asStructType body_t
  rettype' <- asStructType rettype
  pure $
    AppExp
      ( LetFun
          name
          (tparams', params', retdecl', Info (RetType [] rettype'), e')
          body'
          loc
      )
      (Info $ AppRes body_t' [])
--
checkExp (AppExp (Range ExpBase NoInfo TyVar
start Maybe (ExpBase NoInfo TyVar)
maybe_step Inclusiveness (ExpBase NoInfo TyVar)
end SrcLoc
loc) NoInfo AppRes
_) = do
  start' <- Text -> [PrimType] -> Exp -> TermM Exp
require Text
"use in range expression" [PrimType]
anyIntType (Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
start
  let check ExpBase NoInfo TyVar
e = do
        e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
        start_t <- expType start'
        e_t <- expType e'
        ctEq (Reason (locOf e')) start_t e_t
        pure e'
  maybe_step' <- traverse check maybe_step
  end' <- traverse check end
  range_t <- newElemType loc "range" NoMode
  range_t' <- asType range_t
  start_t <- expType start'
  ctEq (Reason (locOf start')) range_t' (arrayOfRank 1 start_t)
  pure $ AppExp (Range start' maybe_step' end' loc) $ Info $ AppRes range_t []
--
checkExp (Project Name
k ExpBase NoInfo TyVar
e NoInfo StructType
NoInfo SrcLoc
loc) = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  kt <- newType loc Lifted "kt" NoMode
  t <- newTypeWithField loc "t" k kt
  e_t <- expType e'
  ctEq (Reason (locOf e')) e_t t
  kt' <- asStructType kt
  pure $ Project k e' (Info kt') loc
--
checkExp (Update ExpBase NoInfo TyVar
src [UpdateStep NoInfo TyVar]
steps ExpBase NoInfo TyVar
ve NoInfo StructType
NoInfo SrcLoc
loc) = do
  src' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
src
  src_t <- expType src'
  src_t' <- asStructType src_t
  ve' <- checkExp ve
  ve_t <- expType ve'
  steps' <- mapM checkStep steps
  mustHaveSteps src' src_t steps' ve_t
  pure $ Update src' steps' ve' (Info src_t') loc

--
checkExp (AppExp (Index ExpBase NoInfo TyVar
e SliceBase NoInfo TyVar
slice SrcLoc
loc) NoInfo AppRes
_) = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  e_t <- expType e'
  slice' <- checkSlice slice
  index_tv <- newTyVar loc Unlifted "index"
  index_elem_t <- newElemType loc "index_elem" NoMode
  let num_slices = SliceBase NoInfo TyVar -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SliceBase NoInfo TyVar -> Int) -> SliceBase NoInfo TyVar -> Int
forall a b. (a -> b) -> a -> b
$ (DimIndexBase NoInfo TyVar -> Bool)
-> SliceBase NoInfo TyVar -> SliceBase NoInfo TyVar
forall a. (a -> Bool) -> [a] -> [a]
filter DimIndexBase NoInfo TyVar -> Bool
forall (f :: * -> *) vn. DimIndexBase f vn -> Bool
isSlice SliceBase NoInfo TyVar
slice
  ctEq (Reason (locOf loc)) (tyVarType NoMode index_tv) $ arrayOfRank num_slices index_elem_t
  ctEq (Reason (locOf e')) e_t $ arrayOfRank (length slice) index_elem_t
  pure $ AppExp (Index e' slice' loc) (Info $ AppRes (tyVarType NoMode index_tv) [])
--
checkExp (AppExp (LetWith IdentBase NoInfo TyVar StructType
dest IdentBase NoInfo TyVar StructType
src [UpdateStep NoInfo TyVar]
steps ExpBase NoInfo TyVar
ve ExpBase NoInfo TyVar
body SrcLoc
loc) NoInfo AppRes
_) = do
  src_t <- SrcLoc -> QualName TyVar -> TermM (TypeBase () NoMode)
lookupVar (IdentBase NoInfo TyVar StructType -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf IdentBase NoInfo TyVar StructType
src) (QualName TyVar -> TermM (TypeBase () NoMode))
-> QualName TyVar -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ TyVar -> QualName TyVar
forall v. v -> QualName v
qualName (TyVar -> QualName TyVar) -> TyVar -> QualName TyVar
forall a b. (a -> b) -> a -> b
$ IdentBase NoInfo TyVar StructType -> TyVar
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase NoInfo TyVar StructType
src
  src_t' <- asStructType src_t
  let src' = IdentBase NoInfo TyVar StructType
src {identType = Info src_t'}
      dest' = IdentBase NoInfo TyVar StructType
dest {identType = Info src_t'}
  steps' <- mapM checkStep steps
  ve' <- checkExp ve
  ve_t <- expType ve'
  mustHaveSteps src' src_t steps' ve_t
  bind [dest'] $ do
    body' <- checkExp body
    body_t <- expType body'
    body_t' <- asStructType body_t
    pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t' [])
--
checkExp (AppExp (If ExpBase NoInfo TyVar
e1 ExpBase NoInfo TyVar
e2 ExpBase NoInfo TyVar
e3 SrcLoc
loc) NoInfo AppRes
_) = do
  e1' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e1
  e1_t <- expType e1'
  e2' <- checkExp e2
  e2_t <- expType e2'
  e3' <- checkExp e3
  e3_t <- expType e3'
  if_t <- newType loc SizeLifted "if_t" NoMode

  ctEq (Reason (locOf e1')) e1_t (Scalar (Prim Bool))
  ctEq (ReasonBranches (locOf loc) e2_t e3_t) e2_t if_t
  ctEq (ReasonBranches (locOf loc) e2_t e3_t) e3_t if_t

  if_t' <- asStructType if_t
  pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes if_t' [])
--
checkExp (AppExp (Match ExpBase NoInfo TyVar
e NonEmpty (CaseBase NoInfo TyVar)
cs SrcLoc
loc) NoInfo AppRes
_) = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  e_t <- expType e'
  (cs', t) <- checkCases e_t cs

  match_t <- newType loc SizeLifted "match_t" NoMode
  ctEq (Reason (locOf loc)) match_t t

  match_t' <- asStructType match_t
  pure $ AppExp (Match e' cs' loc) (Info $ AppRes match_t' [])
--
checkExp (AppExp (Loop [TyVar]
_ PatBase NoInfo TyVar ParamType
pat LoopInitBase NoInfo TyVar
arg LoopFormBase NoInfo TyVar
form ExpBase NoInfo TyVar
body SrcLoc
loc) NoInfo AppRes
_) = do
  arg' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp (ExpBase NoInfo TyVar -> TermM Exp)
-> ExpBase NoInfo TyVar -> TermM Exp
forall a b. (a -> b) -> a -> b
$ case LoopInitBase NoInfo TyVar
arg of
    LoopInitExplicit ExpBase NoInfo TyVar
e -> ExpBase NoInfo TyVar
e
    LoopInitImplicit NoInfo (ExpBase NoInfo TyVar)
_ ->
      -- Should have been filled out in Names
      [Char] -> ExpBase NoInfo TyVar
forall a. HasCallStack => [Char] -> a
error [Char]
"Unspected LoopInitImplicit"
  arg_t <- expType arg'
  loop_t <- newType loc SizeLifted "loop_t" NoMode
  ctEq (Reason (locOf loc)) arg_t loop_t
  bindLetPat pat arg_t $ \Pat ParamType
pat' -> do
    (form', body') <-
      case LoopFormBase NoInfo TyVar
form of
        For (Ident TyVar
i NoInfo StructType
_ SrcLoc
iloc) ExpBase NoInfo TyVar
bound -> do
          bound' <- Text -> [PrimType] -> Exp -> TermM Exp
require Text
"loop bound" [PrimType]
anyIntType (Exp -> TermM Exp) -> TermM Exp -> TermM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
bound
          bound_t <- expType bound'
          bound_t' <- asStructType bound_t
          let i' = TyVar
-> Info StructType -> SrcLoc -> IdentBase Info TyVar StructType
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident TyVar
i (StructType -> Info StructType
forall a. a -> Info a
Info StructType
bound_t') SrcLoc
iloc
          bind [i'] $ do
            body' <- checkExp body
            pure (For i' bound', body')
        While ExpBase NoInfo TyVar
cond -> do
          cond' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
cond
          body' <- checkExp body
          pure (While cond', body')
        ForIn PatBase NoInfo TyVar StructType
elemp ExpBase NoInfo TyVar
arr -> do
          arr' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
arr
          elem_t <- newElemType elemp "elem" NoMode
          arr_t <- expType arr'
          elem_t' <- asType elem_t
          ctEq (Reason (locOf arr')) arr_t $ arrayOfRank 1 elem_t'
          bindLetPat elemp elem_t' $ \Pat ParamType
elemp' -> do
            body' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
body
            pure (ForIn (toStruct <$> elemp') arr', body')
    body_t <- expType body'
    ctEq (Reason (locOf loc)) arg_t body_t
    pure $
      AppExp
        (Loop [] pat' (LoopInitExplicit arg') form' body' loc)
        (Info (AppRes (patternStructType pat') []))
--
checkExp (Ascript ExpBase NoInfo TyVar
e TypeExp (ExpBase NoInfo TyVar) TyVar
te SrcLoc
loc) = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  (te', _, RetType _ st, _) <- checkTypeExp checkSizeExp' te
  e_t <- expType e'
  st' <- asType st
  ctEq (ReasonAscription (locOf e') (toStruct st') (toStruct e_t)) e_t st'
  pure $ Ascript e' te' loc
checkExp (Coerce ExpBase NoInfo TyVar
e TypeExp (ExpBase NoInfo TyVar) TyVar
te NoInfo StructType
NoInfo SrcLoc
loc) = do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  (te', _, RetType _ st, _) <- checkTypeExp checkSizeExp' te
  e_t <- expType e'
  st' <- asType st
  ctEq (Reason (locOf e')) e_t st'
  pure $ Coerce e' te' (Info (toStruct st)) loc

doDefault ::
  [VName] ->
  VName ->
  Either [PrimType] (TypeBase () NoMode) ->
  TermM (TypeBase () NoMode)
doDefault :: [TyVar]
-> TyVar
-> Either [PrimType] (TypeBase () NoMode)
-> TermM (TypeBase () NoMode)
doDefault [TyVar]
tyvars_at_toplevel TyVar
v (Left [PrimType]
pts)
  | [PrimType
pt] <- [PrimType]
pts =
      TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase () NoMode -> TermM (TypeBase () NoMode))
-> TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase () NoMode -> TypeBase () NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase () NoMode -> TypeBase () NoMode)
-> ScalarTypeBase () NoMode -> TypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase () NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
pt
  | IntType -> PrimType
Signed IntType
Int32 PrimType -> [PrimType] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PrimType]
pts = do
      Bool -> TermM () -> TermM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (TyVar
v TyVar -> [TyVar] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TyVar]
tyvars_at_toplevel) (TermM () -> TermM ()) -> TermM () -> TermM ()
forall a b. (a -> b) -> a -> b
$
        Usage -> Doc () -> TermM ()
forall loc. Located loc => loc -> Doc () -> TermM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn Usage
usage Doc ()
"Defaulting ambiguous type to i32."
      TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase () NoMode -> TermM (TypeBase () NoMode))
-> TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase () NoMode -> TypeBase () NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase () NoMode -> TypeBase () NoMode)
-> ScalarTypeBase () NoMode -> TypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase () NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase () NoMode)
-> PrimType -> ScalarTypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int32
  | FloatType -> PrimType
FloatType FloatType
Float64 PrimType -> [PrimType] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PrimType]
pts = do
      Bool -> TermM () -> TermM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (TyVar
v TyVar -> [TyVar] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TyVar]
tyvars_at_toplevel) (TermM () -> TermM ()) -> TermM () -> TermM ()
forall a b. (a -> b) -> a -> b
$
        Usage -> Doc () -> TermM ()
forall loc. Located loc => loc -> Doc () -> TermM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc () -> m ()
warn Usage
usage Doc ()
"Defaulting ambiguous type to f64."
      TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase () NoMode -> TermM (TypeBase () NoMode))
-> TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase () NoMode -> TypeBase () NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase () NoMode -> TypeBase () NoMode)
-> ScalarTypeBase () NoMode -> TypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase () NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase () NoMode)
-> PrimType -> ScalarTypeBase () NoMode
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
FloatType FloatType
Float64
  | Bool
otherwise =
      Usage -> Notes -> Doc () -> TermM (TypeBase () NoMode)
forall loc a. Located loc => loc -> Notes -> Doc () -> TermM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Usage
usage Notes
forall a. Monoid a => a
mempty (Doc () -> TermM (TypeBase () NoMode))
-> (Doc () -> Doc ()) -> Doc () -> TermM (TypeBase () NoMode)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
withIndexLink Doc ()
"ambiguous-type" (Doc () -> TermM (TypeBase () NoMode))
-> Doc () -> TermM (TypeBase () NoMode)
forall a b. (a -> b) -> a -> b
$
        Doc ()
"Type is ambiguous (could be one of"
          Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> [Doc ()] -> Doc ()
forall a. [Doc a] -> Doc a
commasep ((PrimType -> Doc ()) -> [PrimType] -> [Doc ()]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. PrimType -> Doc ann
pretty [PrimType]
pts)
          Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
")."
            Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"Add a type annotation to disambiguate the type."
  where
    usage :: Usage
usage = Loc -> Text -> Usage
forall a. Located a => a -> Text -> Usage
mkUsage Loc
NoLoc Text
"overload"
doDefault [TyVar]
_ TyVar
_ (Right TypeBase () NoMode
t) = TypeBase () NoMode -> TermM (TypeBase () NoMode)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase () NoMode
t

-- | Apply defaults on otherwise ambiguous types. This may result in
-- some type variables becoming known, so we have to perform
-- substitutions on the RHS of the substitutions afterwards.
doDefaults ::
  [VName] ->
  M.Map TyVar (Either [PrimType] (TypeBase () NoMode)) ->
  TermM (M.Map TyVar (TypeBase () NoMode))
doDefaults :: [TyVar]
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
doDefaults [TyVar]
tyvars_at_toplevel Map TyVar (Either [PrimType] (TypeBase () NoMode))
substs = do
  substs' <- (TyVar
 -> Either [PrimType] (TypeBase () NoMode)
 -> TermM (TypeBase () NoMode))
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
M.traverseWithKey ([TyVar]
-> TyVar
-> Either [PrimType] (TypeBase () NoMode)
-> TermM (TypeBase () NoMode)
doDefault [TyVar]
tyvars_at_toplevel) Map TyVar (Either [PrimType] (TypeBase () NoMode))
substs
  pure $ M.map (substTyVars (`M.lookup` substs')) substs'

generalise ::
  TypeBase () NoMode ->
  [UnconTyVar] ->
  Solution ->
  ([TypeParam], [VName])
generalise :: TypeBase () NoMode
-> [UnconTyVar]
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> ([TypeParam], [TyVar])
generalise TypeBase () NoMode
fun_t [UnconTyVar]
unconstrained Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution =
  -- Candidates for let-generalisation are those type variables that
  -- are used in fun_t.
  let visible :: Set TyVar
visible = (TyVar -> Set TyVar) -> Set TyVar -> Set TyVar
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TyVar -> Set TyVar
expandTyVars (Set TyVar -> Set TyVar) -> Set TyVar -> Set TyVar
forall a b. (a -> b) -> a -> b
$ TypeBase () NoMode -> Set TyVar
forall dim o. TypeBase dim o -> Set TyVar
typeVars TypeBase () NoMode
fun_t
      onTyVar :: UnconTyVar -> Either TypeParam TyVar
onTyVar (TyVar
v, Liftedness
l)
        | TyVar
v TyVar -> Set TyVar -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set TyVar
visible = TypeParam -> Either TypeParam TyVar
forall a b. a -> Either a b
Left (TypeParam -> Either TypeParam TyVar)
-> TypeParam -> Either TypeParam TyVar
forall a b. (a -> b) -> a -> b
$ Liftedness -> TyVar -> SrcLoc -> TypeParam
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
l TyVar
v SrcLoc
forall a. Monoid a => a
mempty
        | Bool
otherwise = TyVar -> Either TypeParam TyVar
forall a b. b -> Either a b
Right TyVar
v
   in [Either TypeParam TyVar] -> ([TypeParam], [TyVar])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either TypeParam TyVar] -> ([TypeParam], [TyVar]))
-> [Either TypeParam TyVar] -> ([TypeParam], [TyVar])
forall a b. (a -> b) -> a -> b
$ (UnconTyVar -> Either TypeParam TyVar)
-> [UnconTyVar] -> [Either TypeParam TyVar]
forall a b. (a -> b) -> [a] -> [b]
map UnconTyVar -> Either TypeParam TyVar
onTyVar [UnconTyVar]
unconstrained
  where
    expandTyVars :: TyVar -> Set TyVar
expandTyVars TyVar
v =
      case TyVar
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> Maybe (Either [PrimType] (TypeBase () NoMode))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TyVar
v Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution of
        Just (Right TypeBase () NoMode
t) -> (TyVar -> Set TyVar) -> Set TyVar -> Set TyVar
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TyVar -> Set TyVar
expandTyVars (Set TyVar -> Set TyVar) -> Set TyVar -> Set TyVar
forall a b. (a -> b) -> a -> b
$ TypeBase () NoMode -> Set TyVar
forall dim o. TypeBase dim o -> Set TyVar
typeVars TypeBase () NoMode
t
        Maybe (Either [PrimType] (TypeBase () NoMode))
_ -> TyVar -> Set TyVar
forall a. a -> Set a
S.singleton TyVar
v

generaliseAndDefaults ::
  [UnconTyVar] ->
  Solution ->
  TypeBase () NoMode ->
  TermM ([TypeParam], M.Map VName (TypeBase () NoMode))
generaliseAndDefaults :: [UnconTyVar]
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> TypeBase () NoMode
-> TermM ([TypeParam], Map TyVar (TypeBase () NoMode))
generaliseAndDefaults [UnconTyVar]
unconstrained Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution TypeBase () NoMode
t = do
  let ([TypeParam]
generalised, [TyVar]
unconstrained') =
        TypeBase () NoMode
-> [UnconTyVar]
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> ([TypeParam], [TyVar])
generalise TypeBase () NoMode
t [UnconTyVar]
unconstrained Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution
      -- See #1552 for why we resolve unconstrained and un-generalised type
      -- variables to ().
      units :: Map TyVar (Either a (TypeBase dim o))
units = [(TyVar, Either a (TypeBase dim o))]
-> Map TyVar (Either a (TypeBase dim o))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ((TyVar -> (TyVar, Either a (TypeBase dim o)))
-> [TyVar] -> [(TyVar, Either a (TypeBase dim o))]
forall a b. (a -> b) -> [a] -> [b]
map (,TypeBase dim o -> Either a (TypeBase dim o)
forall a b. b -> Either a b
Right (ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Map Name (TypeBase dim o) -> ScalarTypeBase dim o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record Map Name (TypeBase dim o)
forall a. Monoid a => a
mempty))) [TyVar]
unconstrained')
  solution' <- [TyVar]
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
doDefaults (Set TyVar -> [TyVar]
forall a. Set a -> [a]
S.toList (Set TyVar -> [TyVar]) -> Set TyVar -> [TyVar]
forall a b. (a -> b) -> a -> b
$ TypeBase () NoMode -> Set TyVar
forall dim o. TypeBase dim o -> Set TyVar
typeVars TypeBase () NoMode
t) (Map TyVar (Either [PrimType] (TypeBase () NoMode))
forall {a} {dim} {o}. Map TyVar (Either a (TypeBase dim o))
units Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
forall a. Semigroup a => a -> a -> a
<> Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution)
  pure
    ( generalised,
      solution'
    )

-- | Verify that the recorded type parameter instantiations respect the
-- liftedness of the type parameters. The constraint solver merely propagates
-- liftedness constraints; this is where they are enforced for instantiations,
-- as only here do we know why the constraints exist. (Other liftedness rules
-- are enforced by 'localChecks' in Language.Futhark.TypeChecker.Terms.)
checkTyInstLiftedness :: Solution -> TermM ()
checkTyInstLiftedness :: Map TyVar (Either [PrimType] (TypeBase () NoMode)) -> TermM ()
checkTyInstLiftedness Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution = do
  typarams <- (TermState -> TyParams) -> TermM TyParams
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TermState -> TyParams
termTyParams
  tyset <- asks termTySet
  mapM_ (check typarams tyset) . reverse =<< gets termTyInsts
  where
    -- A Lifted type parameter permits any instantiation.
    check :: Map TyVar (a, Liftedness, a) -> TySet -> TyInst -> f ()
check Map TyVar (a, Liftedness, a)
_ TySet
_ (TyInst Loc
_ QualName TyVar
_ Liftedness
Lifted TyVar
_) = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    check Map TyVar (a, Liftedness, a)
typarams TySet
tyset (TyInst Loc
loc QualName TyVar
qn Liftedness
l TyVar
v)
      | Just (Right TypeBase () NoMode
t) <- TyVar
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> Maybe (Either [PrimType] (TypeBase () NoMode))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TyVar
v Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution = do
          Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (TypeBase () NoMode -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero TypeBase () NoMode
t) (f () -> f ()) -> (Doc () -> f ()) -> Doc () -> f ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Loc -> Notes -> Doc () -> f ()
forall loc a. Located loc => loc -> Notes -> Doc () -> f a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> f ()) -> Doc () -> f ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Type"
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (TypeBase () NoMode -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase () NoMode -> Doc ann
pretty TypeBase () NoMode
t)
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
"found to be functional."
              Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
forall {ann}. Doc ann
when_inst
          let bad :: [Liftedness]
bad = case Liftedness
l of
                Liftedness
Unlifted -> [Liftedness
Lifted, Liftedness
SizeLifted]
                Liftedness
_ -> [Liftedness
Lifted]
              -- The liftedness of a type variable is given by its
              -- binding if it is a type parameter, and by the type
              -- set if it is an abstract type.
              badVar :: QualName TyVar -> Maybe (Doc ann)
badVar QualName TyVar
qv =
                case TyVar -> Map TyVar (a, Liftedness, a) -> Maybe (a, Liftedness, a)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName TyVar -> TyVar
forall vn. QualName vn -> vn
qualLeaf QualName TyVar
qv) Map TyVar (a, Liftedness, a)
typarams of
                  Just (a
_, Liftedness
pl, a
ploc) -> do
                    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Liftedness
pl Liftedness -> [Liftedness] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Liftedness]
bad
                    Doc ann -> Maybe (Doc ann)
forall a. a -> Maybe a
Just (Doc ann -> Maybe (Doc ann)) -> Doc ann -> Maybe (Doc ann)
forall a b. (a -> b) -> a -> b
$
                      Doc ann
"Type parameter"
                        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
dquotes (TyVar -> Doc ann
forall a. TyVar -> Doc a
forall v a. IsName v => v -> Doc a
prettyName (QualName TyVar -> TyVar
forall vn. QualName vn -> vn
qualLeaf QualName TyVar
qv))
                        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann
"bound at"
                        Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> [Char] -> Doc ann
forall ann. [Char] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (a -> [Char]
forall a. Located a => a -> [Char]
locStr a
ploc)
                  Maybe (a, Liftedness, a)
Nothing -> do
                    al <- QualName TyVar -> TySet -> Maybe Liftedness
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup QualName TyVar
qv TySet
tyset
                    guard $ al `elem` bad
                    Just $ "Type" <+> dquotes (pretty qv)
          case (QualName TyVar -> Maybe (Doc ())) -> [QualName TyVar] -> [Doc ()]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe QualName TyVar -> Maybe (Doc ())
forall {ann}. QualName TyVar -> Maybe (Doc ann)
badVar ([QualName TyVar] -> [Doc ()]) -> [QualName TyVar] -> [Doc ()]
forall a b. (a -> b) -> a -> b
$ TypeBase () NoMode -> [QualName TyVar]
forall dim o. TypeBase dim o -> [QualName TyVar]
typeQualVars TypeBase () NoMode
t of
            Doc ()
what : [Doc ()]
_ ->
              Loc -> Notes -> Doc () -> f ()
forall loc a. Located loc => loc -> Notes -> Doc () -> f a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc () -> m a
typeError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> f ()) -> Doc () -> f ()
forall a b. (a -> b) -> a -> b
$
                Doc ()
what
                  Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
<+> case Liftedness
l of
                    Liftedness
Unlifted -> Doc ()
"is lifted and cannot be an array element."
                    Liftedness
_ -> Doc ()
"is lifted and may be a functional type."
                  Doc () -> Doc () -> Doc ()
forall a. Doc a -> Doc a -> Doc a
</> Doc ()
forall {ann}. Doc ann
when_inst
            [] -> () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      where
        when_inst :: Doc ann
when_inst =
          Doc ann
"When instantiating type parameter of" Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
dquotes (QualName TyVar -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. QualName TyVar -> Doc ann
pretty QualName TyVar
qn) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"."

-- | Check a potentially recursive function body. The function is bound to a
-- monomorphic type with its actual parameters, which are known, and a fresh
-- type variable as its return type; that variable is then constrained to the
-- type of the body, and the constraint solver ties the knot. The parameters must
-- be known up front, as unification would otherwise give the function the diets
-- of its first recursive application. A parameterless binding cannot be
-- recursive (see 'resolveValBind'), so it is checked with no self-reference in
-- scope. See Note [Checking recursive functions] in
-- Language.Futhark.TypeChecker.Terms.
checkRecursive ::
  VName ->
  SrcLoc ->
  [Pat ParamType] ->
  ExpBase NoInfo VName ->
  TermM (ExpBase Info VName)
checkRecursive :: TyVar
-> SrcLoc -> [Pat ParamType] -> ExpBase NoInfo TyVar -> TermM Exp
checkRecursive TyVar
_ SrcLoc
_ [] ExpBase NoInfo TyVar
body = ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
body
checkRecursive TyVar
fname SrcLoc
loc [Pat ParamType]
params' ExpBase NoInfo TyVar
body = do
  ret <- SrcLoc
-> Liftedness -> Name -> Freshness -> TermM (TypeBase () Freshness)
forall loc o dim.
Located loc =>
loc -> Liftedness -> Name -> o -> TermM (TypeBase dim o)
newType SrcLoc
loc Liftedness
Lifted (TyVar -> Name
baseName TyVar
fname) Freshness
Nonfresh
  let ftype = [TypeBase () Diet]
-> RetTypeBase () Freshness -> TypeBase () NoMode
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType ((Pat ParamType -> TypeBase () Diet)
-> [Pat ParamType] -> [TypeBase () Diet]
forall a b. (a -> b) -> [a] -> [b]
map ((Exp -> ()) -> ParamType -> TypeBase () Diet
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> Exp -> ()
forall a b. a -> b -> a
const ()) (ParamType -> TypeBase () Diet)
-> (Pat ParamType -> ParamType)
-> Pat ParamType
-> TypeBase () Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> ParamType
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType) [Pat ParamType]
params') ([TyVar] -> TypeBase () Freshness -> RetTypeBase () Freshness
forall dim o. [TyVar] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase () Freshness
ret)
      bindF TermScope ()
scope =
        TermScope ()
scope {scopeVtable = M.insert fname (BoundV [] ftype) $ scopeVtable scope}
  body' <- localScope bindF $ checkExp body
  body_t <- expType body'
  ctEq (Reason (locOf loc)) (toStruct ret) (bimap (const ()) (const NoMode) body_t)
  pure body'

-- | Replace artificial variables with the types they denote, so that no
-- artificial variable leaks into the result.
onArtificial ::
  M.Map TyVar (TypeBase () NoMode) ->
  M.Map TyVar (TypeBase () NoMode) ->
  M.Map TyVar (TypeBase () NoMode)
onArtificial :: Map TyVar (TypeBase () NoMode)
-> Map TyVar (TypeBase () NoMode) -> Map TyVar (TypeBase () NoMode)
onArtificial Map TyVar (TypeBase () NoMode)
artificial Map TyVar (TypeBase () NoMode)
solution =
  (TypeBase () NoMode -> TypeBase () NoMode)
-> Map TyVar (TypeBase () NoMode) -> Map TyVar (TypeBase () NoMode)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TyVar -> Maybe (TypeBase () NoMode))
-> TypeBase () NoMode -> TypeBase () NoMode
forall o d.
Monoid o =>
(TyVar -> Maybe (TypeBase d NoMode))
-> TypeBase d o -> TypeBase d o
substTyVars (TyVar
-> Map TyVar (TypeBase () NoMode) -> Maybe (TypeBase () NoMode)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map TyVar (TypeBase () NoMode)
solution) (TypeBase () NoMode -> TypeBase () NoMode)
-> (TypeBase () NoMode -> TypeBase () NoMode)
-> TypeBase () NoMode
-> TypeBase () NoMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (() -> ()) -> TypeBase () NoMode -> TypeBase () NoMode
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> () -> ()
forall a b. a -> b -> a
const ())) Map TyVar (TypeBase () NoMode)
artificial
    Map TyVar (TypeBase () NoMode)
-> Map TyVar (TypeBase () NoMode) -> Map TyVar (TypeBase () NoMode)
forall a. Semigroup a => a -> a -> a
<> Map TyVar (TypeBase () NoMode)
solution

-- | Type check a single value definition.
checkValDef ::
  ( VName,
    Maybe (TypeExp (ExpBase NoInfo VName) VName),
    [TypeParam],
    [PatBase NoInfo VName ParamType],
    ExpBase NoInfo VName,
    SrcLoc
  ) ->
  TypeM
    ( Either TypeError ([TypeParam], M.Map TyVar (TypeBase () NoMode)),
      [Pat ParamType],
      Maybe (TypeExp Exp VName),
      Exp
    )
checkValDef :: (TyVar, Maybe (TypeExp (ExpBase NoInfo TyVar) TyVar), [TypeParam],
 [PatBase NoInfo TyVar ParamType], ExpBase NoInfo TyVar, SrcLoc)
-> TypeM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp)
checkValDef (TyVar
fname, Maybe (TypeExp (ExpBase NoInfo TyVar) TyVar)
retdecl, [TypeParam]
tparams, [PatBase NoInfo TyVar ParamType]
params, ExpBase NoInfo TyVar
body, SrcLoc
loc) = TermM
  (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
   [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp)
-> TypeM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp)
forall a. TermM a -> TypeM a
runTermM (TermM
   (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
    [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp)
 -> TypeM
      (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
       [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp))
-> TermM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp)
-> TypeM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      [Pat ParamType], Maybe (TypeExp Exp TyVar), Exp)
forall a b. (a -> b) -> a -> b
$ do
  (params', body', retdecl') <-
    [TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType]
    -> TermM ([Pat ParamType], Exp, Maybe (TypeExp Exp TyVar)))
-> TermM ([Pat ParamType], Exp, Maybe (TypeExp Exp TyVar))
forall a.
[TypeParam]
-> [PatBase NoInfo TyVar ParamType]
-> ([Pat ParamType] -> TermM a)
-> TermM a
bindParams [TypeParam]
tparams [PatBase NoInfo TyVar ParamType]
params (([Pat ParamType]
  -> TermM ([Pat ParamType], Exp, Maybe (TypeExp Exp TyVar)))
 -> TermM ([Pat ParamType], Exp, Maybe (TypeExp Exp TyVar)))
-> ([Pat ParamType]
    -> TermM ([Pat ParamType], Exp, Maybe (TypeExp Exp TyVar)))
-> TermM ([Pat ParamType], Exp, Maybe (TypeExp Exp TyVar))
forall a b. (a -> b) -> a -> b
$ \[Pat ParamType]
params' -> do
      body' <- TyVar
-> SrcLoc -> [Pat ParamType] -> ExpBase NoInfo TyVar -> TermM Exp
checkRecursive TyVar
fname SrcLoc
loc [Pat ParamType]
params' ExpBase NoInfo TyVar
body
      (_, retdecl') <- checkRetDecl body' retdecl
      pure (params', body', retdecl')

  cts <- gets termConstraints
  tyvars <- gets termTyVars
  typarams <- gets termTyParams
  artificial <- gets $ M.map (first (const ())) . termArtificial

  solution <-
    bitraverse
      pure
      (fmap (second (onArtificial artificial)) . onTySolution params' body')
      $ solve (reverse cts) typarams tyvars
  pure (solution, params', retdecl', body')
  where
    onTySolution :: [Pat (TypeBase a Diet)]
-> Exp
-> ([UnconTyVar],
    Map TyVar (Either [PrimType] (TypeBase () NoMode)))
-> TermM ([TypeParam], Map TyVar (TypeBase () NoMode))
onTySolution [Pat (TypeBase a Diet)]
params' Exp
body' ([UnconTyVar]
unconstrained, Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution) = do
      Map TyVar (Either [PrimType] (TypeBase () NoMode)) -> TermM ()
checkTyInstLiftedness Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution
      body_t <- Exp -> TermM (TypeBase () NoMode)
expType Exp
body'
      let fun_t =
            [TypeBase () Diet]
-> RetTypeBase () Freshness -> TypeBase () NoMode
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType
              ((Pat (TypeBase a Diet) -> TypeBase () Diet)
-> [Pat (TypeBase a Diet)] -> [TypeBase () Diet]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> ()) -> TypeBase a Diet -> TypeBase () Diet
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> a -> ()
forall a b. a -> b -> a
const ()) (TypeBase a Diet -> TypeBase () Diet)
-> (Pat (TypeBase a Diet) -> TypeBase a Diet)
-> Pat (TypeBase a Diet)
-> TypeBase () Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (TypeBase a Diet) -> TypeBase a Diet
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType) [Pat (TypeBase a Diet)]
params')
              ([TyVar] -> TypeBase () Freshness -> RetTypeBase () Freshness
forall dim o. [TyVar] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (TypeBase () Freshness -> RetTypeBase () Freshness)
-> TypeBase () Freshness -> RetTypeBase () Freshness
forall a b. (a -> b) -> a -> b
$ (() -> ())
-> (NoMode -> Freshness)
-> TypeBase () NoMode
-> TypeBase () Freshness
forall a b c d.
(a -> b) -> (c -> d) -> TypeBase a c -> TypeBase b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (() -> () -> ()
forall a b. a -> b -> a
const ()) (Freshness -> NoMode -> Freshness
forall a b. a -> b -> a
const Freshness
Nonfresh) TypeBase () NoMode
body_t)
      generaliseAndDefaults unconstrained solution fun_t

-- | Type check a single expression, which may have a polymorphic
-- type.
checkSingleExp ::
  ExpBase NoInfo VName ->
  TypeM
    ( Either TypeError ([TypeParam], M.Map TyVar (TypeBase () NoMode)),
      Exp
    )
checkSingleExp :: ExpBase NoInfo TyVar
-> TypeM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      Exp)
checkSingleExp ExpBase NoInfo TyVar
e = TermM
  (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
   Exp)
-> TypeM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      Exp)
forall a. TermM a -> TypeM a
runTermM (TermM
   (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
    Exp)
 -> TypeM
      (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
       Exp))
-> TermM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      Exp)
-> TypeM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      Exp)
forall a b. (a -> b) -> a -> b
$ do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkExp ExpBase NoInfo TyVar
e
  cts <- gets termConstraints
  tyvars <- gets termTyVars
  typarams <- gets termTyParams
  artificial <- gets termArtificial

  case solve cts typarams tyvars of
    Left TypeError
err -> (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
 Exp)
-> TermM
     (Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode)),
      Exp)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeError
-> Either TypeError ([TypeParam], Map TyVar (TypeBase () NoMode))
forall a b. a -> Either a b
Left TypeError
err, Exp
e')
    Right ([UnconTyVar]
unconstrained, Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution) -> do
      Map TyVar (Either [PrimType] (TypeBase () NoMode)) -> TermM ()
checkTyInstLiftedness Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution
      e_t <- Exp -> TermM (TypeBase () NoMode)
expType Exp
e'
      x <-
        second (onArtificial (M.map (first (const ())) artificial))
          <$> generaliseAndDefaults unconstrained solution (first (const ()) e_t)
      pure (Right x, e')

-- | Type-check a single size expression in isolation, which must have
-- type @i64@.
checkSizeExp ::
  ExpBase NoInfo VName ->
  TypeM
    ( Either TypeError ([UnconTyVar], M.Map TyVar (TypeBase () NoMode)),
      Exp
    )
checkSizeExp :: ExpBase NoInfo TyVar
-> TypeM
     (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
      Exp)
checkSizeExp ExpBase NoInfo TyVar
e = TermM
  (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
   Exp)
-> TypeM
     (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
      Exp)
forall a. TermM a -> TypeM a
runTermM (TermM
   (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
    Exp)
 -> TypeM
      (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
       Exp))
-> TermM
     (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
      Exp)
-> TypeM
     (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
      Exp)
forall a b. (a -> b) -> a -> b
$ do
  e' <- ExpBase NoInfo TyVar -> TermM Exp
checkSizeExp' ExpBase NoInfo TyVar
e
  cts <- gets termConstraints
  tyvars <- gets termTyVars
  typarams <- gets termTyParams
  artificial <- gets termArtificial

  case solve cts typarams tyvars of
    Left TypeError
err -> (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
 Exp)
-> TermM
     (Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode)),
      Exp)
forall a. a -> TermM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeError
-> Either TypeError ([UnconTyVar], Map TyVar (TypeBase () NoMode))
forall a b. a -> Either a b
Left TypeError
err, Exp
e')
    Right ([UnconTyVar]
unconstrained, Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution) -> do
      Map TyVar (Either [PrimType] (TypeBase () NoMode)) -> TermM ()
checkTyInstLiftedness Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution
      solution' <-
        Map TyVar (TypeBase () NoMode)
-> Map TyVar (TypeBase () NoMode) -> Map TyVar (TypeBase () NoMode)
onArtificial ((TypeBase () NoMode -> TypeBase () NoMode)
-> Map TyVar (TypeBase () NoMode) -> Map TyVar (TypeBase () NoMode)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((() -> ()) -> TypeBase () NoMode -> TypeBase () NoMode
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> () -> ()
forall a b. a -> b -> a
const ())) Map TyVar (TypeBase () NoMode)
artificial)
          (Map TyVar (TypeBase () NoMode) -> Map TyVar (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TyVar]
-> Map TyVar (Either [PrimType] (TypeBase () NoMode))
-> TermM (Map TyVar (TypeBase () NoMode))
doDefaults [TyVar]
forall a. Monoid a => a
mempty Map TyVar (Either [PrimType] (TypeBase () NoMode))
solution
      pure (Right (unconstrained, solution'), e')