{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Trustworthy #-}

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

    -- * Primitive checking
    unifies,
    require,
    checkTypeExpNonrigid,
    checkTypeExpRigid,

    -- * Sizes
    isInt64,
    maybeDimFromExp,
    dimFromExp,
    dimFromArg,
    noSizeEscape,

    -- * Control flow
    collectOccurrences,
    tapOccurrences,
    alternative,
    sequentially,
    incLevel,

    -- * Consumption and uniqueness
    Names,
    Occurrence (..),
    Occurrences,
    onlySelfAliasing,
    noUnique,
    removeSeminullOccurrences,
    occur,
    observe,
    consume,
    consuming,
    observation,
    consumption,
    checkIfConsumable,
    seqOccurrences,
    checkOccurrences,
    allConsumed,

    -- * Errors
    useAfterConsume,
    unusedSize,
    notConsumable,
    unexpectedType,
    uniqueReturnAliased,
    returnAliased,
    badLetWithValue,
    anyConsumption,
    allOccurring,
  )
where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable
import Data.Char (isAscii)
import Data.List (find, isPrefixOf, sort)
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import Futhark.Util.Pretty hiding (bool, group, space)
import Language.Futhark
import Language.Futhark.Semantic (includeToFilePath)
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Monad hiding (BoundV)
import qualified Language.Futhark.TypeChecker.Monad as TypeM
import Language.Futhark.TypeChecker.Types
import Language.Futhark.TypeChecker.Unify hiding (Usage)
import Prelude hiding (mod)

--- Uniqueness

data Usage
  = Consumed SrcLoc
  | Observed SrcLoc
  deriving (Usage -> Usage -> Bool
(Usage -> Usage -> Bool) -> (Usage -> Usage -> Bool) -> Eq Usage
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Usage -> Usage -> Bool
$c/= :: Usage -> Usage -> Bool
== :: Usage -> Usage -> Bool
$c== :: Usage -> Usage -> Bool
Eq, Eq Usage
Eq Usage
-> (Usage -> Usage -> Ordering)
-> (Usage -> Usage -> Bool)
-> (Usage -> Usage -> Bool)
-> (Usage -> Usage -> Bool)
-> (Usage -> Usage -> Bool)
-> (Usage -> Usage -> Usage)
-> (Usage -> Usage -> Usage)
-> Ord Usage
Usage -> Usage -> Bool
Usage -> Usage -> Ordering
Usage -> Usage -> Usage
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Usage -> Usage -> Usage
$cmin :: Usage -> Usage -> Usage
max :: Usage -> Usage -> Usage
$cmax :: Usage -> Usage -> Usage
>= :: Usage -> Usage -> Bool
$c>= :: Usage -> Usage -> Bool
> :: Usage -> Usage -> Bool
$c> :: Usage -> Usage -> Bool
<= :: Usage -> Usage -> Bool
$c<= :: Usage -> Usage -> Bool
< :: Usage -> Usage -> Bool
$c< :: Usage -> Usage -> Bool
compare :: Usage -> Usage -> Ordering
$ccompare :: Usage -> Usage -> Ordering
$cp1Ord :: Eq Usage
Ord, Int -> Usage -> ShowS
[Usage] -> ShowS
Usage -> String
(Int -> Usage -> ShowS)
-> (Usage -> String) -> ([Usage] -> ShowS) -> Show Usage
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Usage] -> ShowS
$cshowList :: [Usage] -> ShowS
show :: Usage -> String
$cshow :: Usage -> String
showsPrec :: Int -> Usage -> ShowS
$cshowsPrec :: Int -> Usage -> ShowS
Show)

type Names = S.Set VName

-- | The consumption set is a Maybe so we can distinguish whether a
-- consumption took place, but the variable went out of scope since,
-- or no consumption at all took place.
data Occurrence = Occurrence
  { Occurrence -> Names
observed :: Names,
    Occurrence -> Maybe Names
consumed :: Maybe Names,
    Occurrence -> SrcLoc
location :: SrcLoc
  }
  deriving (Occurrence -> Occurrence -> Bool
(Occurrence -> Occurrence -> Bool)
-> (Occurrence -> Occurrence -> Bool) -> Eq Occurrence
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Occurrence -> Occurrence -> Bool
$c/= :: Occurrence -> Occurrence -> Bool
== :: Occurrence -> Occurrence -> Bool
$c== :: Occurrence -> Occurrence -> Bool
Eq, Int -> Occurrence -> ShowS
[Occurrence] -> ShowS
Occurrence -> String
(Int -> Occurrence -> ShowS)
-> (Occurrence -> String)
-> ([Occurrence] -> ShowS)
-> Show Occurrence
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Occurrence] -> ShowS
$cshowList :: [Occurrence] -> ShowS
show :: Occurrence -> String
$cshow :: Occurrence -> String
showsPrec :: Int -> Occurrence -> ShowS
$cshowsPrec :: Int -> Occurrence -> ShowS
Show)

instance Located Occurrence where
  locOf :: Occurrence -> Loc
locOf = SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf (SrcLoc -> Loc) -> (Occurrence -> SrcLoc) -> Occurrence -> Loc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurrence -> SrcLoc
location

observation :: Aliasing -> SrcLoc -> Occurrence
observation :: Aliasing -> SrcLoc -> Occurrence
observation = (Names -> Maybe Names -> SrcLoc -> Occurrence)
-> Maybe Names -> Names -> SrcLoc -> Occurrence
forall a b c. (a -> b -> c) -> b -> a -> c
flip Names -> Maybe Names -> SrcLoc -> Occurrence
Occurrence Maybe Names
forall a. Maybe a
Nothing (Names -> SrcLoc -> Occurrence)
-> (Aliasing -> Names) -> Aliasing -> SrcLoc -> Occurrence
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> VName) -> Aliasing -> Names
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar

consumption :: Aliasing -> SrcLoc -> Occurrence
consumption :: Aliasing -> SrcLoc -> Occurrence
consumption = Names -> Maybe Names -> SrcLoc -> Occurrence
Occurrence Names
forall a. Set a
S.empty (Maybe Names -> SrcLoc -> Occurrence)
-> (Aliasing -> Maybe Names) -> Aliasing -> SrcLoc -> Occurrence
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> Maybe Names
forall a. a -> Maybe a
Just (Names -> Maybe Names)
-> (Aliasing -> Names) -> Aliasing -> Maybe Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> VName) -> Aliasing -> Names
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar

-- | A null occurence is one that we can remove without affecting
-- anything.
nullOccurrence :: Occurrence -> Bool
nullOccurrence :: Occurrence -> Bool
nullOccurrence Occurrence
occ = Names -> Bool
forall a. Set a -> Bool
S.null (Occurrence -> Names
observed Occurrence
occ) Bool -> Bool -> Bool
&& Maybe Names -> Bool
forall a. Maybe a -> Bool
isNothing (Occurrence -> Maybe Names
consumed Occurrence
occ)

-- | A seminull occurence is one that does not contain references to
-- any variables in scope.  The big difference is that a seminull
-- occurence may denote a consumption, as long as the array that was
-- consumed is now out of scope.
seminullOccurrence :: Occurrence -> Bool
seminullOccurrence :: Occurrence -> Bool
seminullOccurrence Occurrence
occ = Names -> Bool
forall a. Set a -> Bool
S.null (Occurrence -> Names
observed Occurrence
occ) Bool -> Bool -> Bool
&& Bool -> (Names -> Bool) -> Maybe Names -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True Names -> Bool
forall a. Set a -> Bool
S.null (Occurrence -> Maybe Names
consumed Occurrence
occ)

type Occurrences = [Occurrence]

type UsageMap = M.Map VName [Usage]

usageMap :: Occurrences -> UsageMap
usageMap :: [Occurrence] -> UsageMap
usageMap = (UsageMap -> Occurrence -> UsageMap)
-> UsageMap -> [Occurrence] -> UsageMap
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl UsageMap -> Occurrence -> UsageMap
comb UsageMap
forall k a. Map k a
M.empty
  where
    comb :: UsageMap -> Occurrence -> UsageMap
comb UsageMap
m (Occurrence Names
obs Maybe Names
cons SrcLoc
loc) =
      let m' :: UsageMap
m' = (UsageMap -> VName -> UsageMap) -> UsageMap -> Names -> UsageMap
forall a b. (a -> b -> a) -> a -> Set b -> a
S.foldl' (Usage -> UsageMap -> VName -> UsageMap
forall k a. Ord k => a -> Map k [a] -> k -> Map k [a]
ins (Usage -> UsageMap -> VName -> UsageMap)
-> Usage -> UsageMap -> VName -> UsageMap
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
Observed SrcLoc
loc) UsageMap
m Names
obs
       in (UsageMap -> VName -> UsageMap) -> UsageMap -> Names -> UsageMap
forall a b. (a -> b -> a) -> a -> Set b -> a
S.foldl' (Usage -> UsageMap -> VName -> UsageMap
forall k a. Ord k => a -> Map k [a] -> k -> Map k [a]
ins (Usage -> UsageMap -> VName -> UsageMap)
-> Usage -> UsageMap -> VName -> UsageMap
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
Consumed SrcLoc
loc) UsageMap
m' (Names -> UsageMap) -> Names -> UsageMap
forall a b. (a -> b) -> a -> b
$ Names -> Maybe Names -> Names
forall a. a -> Maybe a -> a
fromMaybe Names
forall a. Monoid a => a
mempty Maybe Names
cons
    ins :: a -> Map k [a] -> k -> Map k [a]
ins a
v Map k [a]
m k
k = ([a] -> [a] -> [a]) -> k -> [a] -> Map k [a] -> Map k [a]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) k
k [a
v] Map k [a]
m

combineOccurrences :: VName -> Usage -> Usage -> TermTypeM Usage
combineOccurrences :: VName -> Usage -> Usage -> TermTypeM Usage
combineOccurrences VName
_ (Observed SrcLoc
loc) (Observed SrcLoc
_) = Usage -> TermTypeM Usage
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Usage -> TermTypeM Usage) -> Usage -> TermTypeM Usage
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
Observed SrcLoc
loc
combineOccurrences VName
name (Consumed SrcLoc
wloc) (Observed SrcLoc
rloc) =
  VName -> SrcLoc -> SrcLoc -> TermTypeM Usage
forall a. VName -> SrcLoc -> SrcLoc -> TermTypeM a
useAfterConsume VName
name SrcLoc
rloc SrcLoc
wloc
combineOccurrences VName
name (Observed SrcLoc
rloc) (Consumed SrcLoc
wloc) =
  VName -> SrcLoc -> SrcLoc -> TermTypeM Usage
forall a. VName -> SrcLoc -> SrcLoc -> TermTypeM a
useAfterConsume VName
name SrcLoc
rloc SrcLoc
wloc
combineOccurrences VName
name (Consumed SrcLoc
loc1) (Consumed SrcLoc
loc2) =
  VName -> SrcLoc -> SrcLoc -> TermTypeM Usage
forall a. VName -> SrcLoc -> SrcLoc -> TermTypeM a
useAfterConsume VName
name (SrcLoc -> SrcLoc -> SrcLoc
forall a. Ord a => a -> a -> a
max SrcLoc
loc1 SrcLoc
loc2) (SrcLoc -> SrcLoc -> SrcLoc
forall a. Ord a => a -> a -> a
min SrcLoc
loc1 SrcLoc
loc2)

checkOccurrences :: Occurrences -> TermTypeM ()
checkOccurrences :: [Occurrence] -> TermTypeM ()
checkOccurrences = TermTypeM (Map VName ()) -> TermTypeM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TermTypeM (Map VName ()) -> TermTypeM ())
-> ([Occurrence] -> TermTypeM (Map VName ()))
-> [Occurrence]
-> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> [Usage] -> TermTypeM ())
-> UsageMap -> TermTypeM (Map VName ())
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
M.traverseWithKey VName -> [Usage] -> TermTypeM ()
comb (UsageMap -> TermTypeM (Map VName ()))
-> ([Occurrence] -> UsageMap)
-> [Occurrence]
-> TermTypeM (Map VName ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Occurrence] -> UsageMap
usageMap
  where
    comb :: VName -> [Usage] -> TermTypeM ()
comb VName
_ [] = () -> TermTypeM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    comb VName
name (Usage
u : [Usage]
us) = (Usage -> Usage -> TermTypeM Usage)
-> Usage -> [Usage] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ (VName -> Usage -> Usage -> TermTypeM Usage
combineOccurrences VName
name) Usage
u [Usage]
us

allObserved :: Occurrences -> Names
allObserved :: [Occurrence] -> Names
allObserved = [Names] -> Names
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ([Names] -> Names)
-> ([Occurrence] -> [Names]) -> [Occurrence] -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Occurrence -> Names) -> [Occurrence] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Occurrence -> Names
observed

allConsumed :: Occurrences -> Names
allConsumed :: [Occurrence] -> Names
allConsumed = [Names] -> Names
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ([Names] -> Names)
-> ([Occurrence] -> [Names]) -> [Occurrence] -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Occurrence -> Names) -> [Occurrence] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Names -> Maybe Names -> Names
forall a. a -> Maybe a -> a
fromMaybe Names
forall a. Monoid a => a
mempty (Maybe Names -> Names)
-> (Occurrence -> Maybe Names) -> Occurrence -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurrence -> Maybe Names
consumed)

allOccurring :: Occurrences -> Names
allOccurring :: [Occurrence] -> Names
allOccurring [Occurrence]
occs = [Occurrence] -> Names
allConsumed [Occurrence]
occs Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Occurrence] -> Names
allObserved [Occurrence]
occs

anyConsumption :: Occurrences -> Maybe Occurrence
anyConsumption :: [Occurrence] -> Maybe Occurrence
anyConsumption = (Occurrence -> Bool) -> [Occurrence] -> Maybe Occurrence
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Maybe Names -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Names -> Bool)
-> (Occurrence -> Maybe Names) -> Occurrence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurrence -> Maybe Names
consumed)

seqOccurrences :: Occurrences -> Occurrences -> Occurrences
seqOccurrences :: [Occurrence] -> [Occurrence] -> [Occurrence]
seqOccurrences [Occurrence]
occurs1 [Occurrence]
occurs2 =
  (Occurrence -> Bool) -> [Occurrence] -> [Occurrence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurrence -> Bool) -> Occurrence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurrence -> Bool
nullOccurrence) ([Occurrence] -> [Occurrence]) -> [Occurrence] -> [Occurrence]
forall a b. (a -> b) -> a -> b
$ (Occurrence -> Occurrence) -> [Occurrence] -> [Occurrence]
forall a b. (a -> b) -> [a] -> [b]
map Occurrence -> Occurrence
filt [Occurrence]
occurs1 [Occurrence] -> [Occurrence] -> [Occurrence]
forall a. [a] -> [a] -> [a]
++ [Occurrence]
occurs2
  where
    filt :: Occurrence -> Occurrence
filt Occurrence
occ =
      Occurrence
occ {observed :: Names
observed = Occurrence -> Names
observed Occurrence
occ Names -> Names -> Names
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Names
postcons}
    postcons :: Names
postcons = [Occurrence] -> Names
allConsumed [Occurrence]
occurs2

altOccurrences :: Occurrences -> Occurrences -> Occurrences
altOccurrences :: [Occurrence] -> [Occurrence] -> [Occurrence]
altOccurrences [Occurrence]
occurs1 [Occurrence]
occurs2 =
  (Occurrence -> Bool) -> [Occurrence] -> [Occurrence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurrence -> Bool) -> Occurrence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurrence -> Bool
nullOccurrence) ([Occurrence] -> [Occurrence]) -> [Occurrence] -> [Occurrence]
forall a b. (a -> b) -> a -> b
$ (Occurrence -> Occurrence) -> [Occurrence] -> [Occurrence]
forall a b. (a -> b) -> [a] -> [b]
map Occurrence -> Occurrence
filt1 [Occurrence]
occurs1 [Occurrence] -> [Occurrence] -> [Occurrence]
forall a. [a] -> [a] -> [a]
++ (Occurrence -> Occurrence) -> [Occurrence] -> [Occurrence]
forall a b. (a -> b) -> [a] -> [b]
map Occurrence -> Occurrence
filt2 [Occurrence]
occurs2
  where
    filt1 :: Occurrence -> Occurrence
filt1 Occurrence
occ =
      Occurrence
occ
        { consumed :: Maybe Names
consumed = Names -> Names -> Names
forall a. Ord a => Set a -> Set a -> Set a
S.difference (Names -> Names -> Names) -> Maybe Names -> Maybe (Names -> Names)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Occurrence -> Maybe Names
consumed Occurrence
occ Maybe (Names -> Names) -> Maybe Names -> Maybe Names
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Names -> Maybe Names
forall (f :: * -> *) a. Applicative f => a -> f a
pure Names
cons2,
          observed :: Names
observed = Occurrence -> Names
observed Occurrence
occ Names -> Names -> Names
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Names
cons2
        }
    filt2 :: Occurrence -> Occurrence
filt2 Occurrence
occ =
      Occurrence
occ
        { consumed :: Maybe Names
consumed = Occurrence -> Maybe Names
consumed Occurrence
occ,
          observed :: Names
observed = Occurrence -> Names
observed Occurrence
occ Names -> Names -> Names
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Names
cons1
        }
    cons1 :: Names
cons1 = [Occurrence] -> Names
allConsumed [Occurrence]
occurs1
    cons2 :: Names
cons2 = [Occurrence] -> Names
allConsumed [Occurrence]
occurs2

-- | Whether something is a global or a local variable.
data Locality = Local | Global
  deriving (Int -> Locality -> ShowS
[Locality] -> ShowS
Locality -> String
(Int -> Locality -> ShowS)
-> (Locality -> String) -> ([Locality] -> ShowS) -> Show Locality
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Locality] -> ShowS
$cshowList :: [Locality] -> ShowS
show :: Locality -> String
$cshow :: Locality -> String
showsPrec :: Int -> Locality -> ShowS
$cshowsPrec :: Int -> Locality -> ShowS
Show)

data ValBinding
  = -- | Aliases in parameters indicate the lexical
    -- closure.
    BoundV Locality [TypeParam] PatType
  | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType)
  | EqualityF
  | WasConsumed SrcLoc
  deriving (Int -> ValBinding -> ShowS
[ValBinding] -> ShowS
ValBinding -> String
(Int -> ValBinding -> ShowS)
-> (ValBinding -> String)
-> ([ValBinding] -> ShowS)
-> Show ValBinding
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ValBinding] -> ShowS
$cshowList :: [ValBinding] -> ShowS
show :: ValBinding -> String
$cshow :: ValBinding -> String
showsPrec :: Int -> ValBinding -> ShowS
$cshowsPrec :: Int -> ValBinding -> ShowS
Show)

--- Errors

describeVar :: SrcLoc -> VName -> TermTypeM Doc
describeVar :: SrcLoc -> VName -> TermTypeM Doc
describeVar SrcLoc
loc VName
v =
  (TermTypeState -> Doc) -> TermTypeM Doc
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((TermTypeState -> Doc) -> TermTypeM Doc)
-> (TermTypeState -> Doc) -> TermTypeM Doc
forall a b. (a -> b) -> a -> b
$
    Doc -> (NameReason -> Doc) -> Maybe NameReason -> Doc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Doc
"variable" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (VName -> Doc
forall v. IsName v => v -> Doc
pprName VName
v)) (SrcLoc -> NameReason -> Doc
nameReason SrcLoc
loc)
      (Maybe NameReason -> Doc)
-> (TermTypeState -> Maybe NameReason) -> TermTypeState -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Map VName NameReason -> Maybe NameReason
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v
      (Map VName NameReason -> Maybe NameReason)
-> (TermTypeState -> Map VName NameReason)
-> TermTypeState
-> Maybe NameReason
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeState -> Map VName NameReason
stateNames

useAfterConsume :: VName -> SrcLoc -> SrcLoc -> TermTypeM a
useAfterConsume :: VName -> SrcLoc -> SrcLoc -> TermTypeM a
useAfterConsume VName
name SrcLoc
rloc SrcLoc
wloc = do
  Doc
name' <- SrcLoc -> VName -> TermTypeM Doc
describeVar SrcLoc
rloc VName
name
  SrcLoc -> Notes -> Doc -> TermTypeM a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
rloc Notes
forall a. Monoid a => a
mempty (Doc -> TermTypeM a) -> (Doc -> Doc) -> Doc -> TermTypeM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc -> Doc
withIndexLink Doc
"use-after-consume" (Doc -> TermTypeM a) -> Doc -> TermTypeM a
forall a b. (a -> b) -> a -> b
$
    Doc
"Using" Doc -> Doc -> Doc
<+> Doc
name' Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
", but this was consumed at"
      Doc -> Doc -> Doc
<+> String -> Doc
text (SrcLoc -> SrcLoc -> String
forall a b. (Located a, Located b) => a -> b -> String
locStrRel SrcLoc
rloc SrcLoc
wloc) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
".  (Possibly through aliasing.)"

badLetWithValue :: (Pretty arr, Pretty src) => arr -> src -> SrcLoc -> TermTypeM a
badLetWithValue :: arr -> src -> SrcLoc -> TermTypeM a
badLetWithValue arr
arre src
vale SrcLoc
loc =
  SrcLoc -> Notes -> Doc -> TermTypeM 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 -> TermTypeM a) -> Doc -> TermTypeM a
forall a b. (a -> b) -> a -> b
$
    Doc
"Source array for in-place update"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (arr -> Doc
forall a. Pretty a => a -> Doc
ppr arr
arre)
      Doc -> Doc -> Doc
</> Doc
"might alias update value"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (src -> Doc
forall a. Pretty a => a -> Doc
ppr src
vale)
      Doc -> Doc -> Doc
</> Doc
"Hint: use" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote Doc
"copy" Doc -> Doc -> Doc
<+> Doc
"to remove aliases from the value."

returnAliased :: Name -> SrcLoc -> TermTypeM ()
returnAliased :: Name -> SrcLoc -> TermTypeM ()
returnAliased Name
name SrcLoc
loc =
  SrcLoc -> Notes -> Doc -> TermTypeM ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> TermTypeM ()) -> (Doc -> Doc) -> Doc -> TermTypeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc -> Doc
withIndexLink Doc
"return-aliased" (Doc -> TermTypeM ()) -> Doc -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
    Doc
"Unique-typed return value is aliased to"
      Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (Name -> Doc
forall v. IsName v => v -> Doc
pprName Name
name) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
", which is not consumable."

uniqueReturnAliased :: SrcLoc -> TermTypeM a
uniqueReturnAliased :: SrcLoc -> TermTypeM a
uniqueReturnAliased SrcLoc
loc =
  SrcLoc -> Notes -> Doc -> TermTypeM 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 -> TermTypeM a) -> (Doc -> Doc) -> Doc -> TermTypeM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc -> Doc
withIndexLink Doc
"unique-return-aliased" (Doc -> TermTypeM a) -> Doc -> TermTypeM a
forall a b. (a -> b) -> a -> b
$
    Doc
"A unique-typed component of the return value is aliased to some other component."

unexpectedType :: MonadTypeChecker m => SrcLoc -> StructType -> [StructType] -> m a
unexpectedType :: SrcLoc -> StructType -> [StructType] -> m a
unexpectedType SrcLoc
loc StructType
_ [] =
  SrcLoc -> Notes -> Doc -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m a) -> Doc -> m a
forall a b. (a -> b) -> a -> b
$
    Doc
"Type of expression at" Doc -> Doc -> Doc
<+> String -> Doc
text (SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc)
      Doc -> Doc -> Doc
<+> Doc
"cannot have any type - possibly a bug in the type checker."
unexpectedType SrcLoc
loc StructType
t [StructType]
ts =
  SrcLoc -> Notes -> Doc -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m a) -> Doc -> m a
forall a b. (a -> b) -> a -> b
$
    Doc
"Type of expression at" Doc -> Doc -> Doc
<+> String -> Doc
text (SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc) Doc -> Doc -> Doc
<+> Doc
"must be one of"
      Doc -> Doc -> Doc
<+> [Doc] -> Doc
commasep ((StructType -> Doc) -> [StructType] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> Doc
forall a. Pretty a => a -> Doc
ppr [StructType]
ts) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
", but is"
      Doc -> Doc -> Doc
<+> StructType -> Doc
forall a. Pretty a => a -> Doc
ppr StructType
t Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."

notConsumable :: MonadTypeChecker m => SrcLoc -> Doc -> m b
notConsumable :: SrcLoc -> Doc -> m b
notConsumable SrcLoc
loc Doc
v =
  SrcLoc -> Notes -> Doc -> m b
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m b) -> (Doc -> Doc) -> Doc -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc -> Doc
withIndexLink Doc
"not-consumable" (Doc -> m b) -> Doc -> m b
forall a b. (a -> b) -> a -> b
$
    Doc
"Would consume" Doc -> Doc -> Doc
<+> Doc
v Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
", which is not consumable."

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

--- Scope management

data InferredType
  = NoneInferred
  | Ascribed PatType

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

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

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

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

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

envToTermScope :: Env -> TermScope
envToTermScope :: Env -> TermScope
envToTermScope Env
env =
  TermScope :: Map VName ValBinding
-> Map VName TypeBinding -> Map VName Mod -> NameMap -> TermScope
TermScope
    { scopeVtable :: Map VName ValBinding
scopeVtable = Map VName ValBinding
vtable,
      scopeTypeTable :: Map VName TypeBinding
scopeTypeTable = Env -> Map VName TypeBinding
envTypeTable Env
env,
      scopeNameMap :: NameMap
scopeNameMap = Env -> NameMap
envNameMap Env
env,
      scopeModTable :: Map VName Mod
scopeModTable = Env -> Map VName Mod
envModTable Env
env
    }
  where
    vtable :: Map VName ValBinding
vtable = (VName -> BoundV -> ValBinding)
-> Map VName BoundV -> Map VName ValBinding
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey VName -> BoundV -> ValBinding
valBinding (Map VName BoundV -> Map VName ValBinding)
-> Map VName BoundV -> Map VName ValBinding
forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable Env
env
    valBinding :: VName -> BoundV -> ValBinding
valBinding VName
k (TypeM.BoundV [TypeParam]
tps StructType
v) =
      Locality -> [TypeParam] -> PatType -> ValBinding
BoundV Locality
Global [TypeParam]
tps (PatType -> ValBinding) -> PatType -> ValBinding
forall a b. (a -> b) -> a -> b
$
        StructType
v
          StructType -> Aliasing -> PatType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` (if StructType -> Int
forall dim as. TypeBase dim as -> Int
arrayRank StructType
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Alias -> Aliasing
forall a. a -> Set a
S.singleton (VName -> Alias
AliasBound VName
k) else Aliasing
forall a. Monoid a => a
mempty)

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

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

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

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

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

-- | A description of where an artificial compiler-generated
-- intermediate name came from.
data NameReason
  = -- | Name is the result of a function application.
    NameAppRes (Maybe (QualName VName)) SrcLoc

nameReason :: SrcLoc -> NameReason -> Doc
nameReason :: SrcLoc -> NameReason -> Doc
nameReason SrcLoc
loc (NameAppRes Maybe (QualName VName)
Nothing SrcLoc
apploc) =
  Doc
"result of application at" Doc -> Doc -> Doc
<+> String -> Doc
text (SrcLoc -> SrcLoc -> String
forall a b. (Located a, Located b) => a -> b -> String
locStrRel SrcLoc
loc SrcLoc
apploc)
nameReason SrcLoc
loc (NameAppRes Maybe (QualName VName)
fname SrcLoc
apploc) =
  Doc
"result of applying" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (Maybe (QualName VName) -> Doc
forall a. Pretty a => a -> Doc
ppr Maybe (QualName VName)
fname)
    Doc -> Doc -> Doc
<+> Doc -> Doc
parens (Doc
"at" Doc -> Doc -> Doc
<+> String -> Doc
text (SrcLoc -> SrcLoc -> String
forall a b. (Located a, Located b) => a -> b -> String
locStrRel SrcLoc
loc SrcLoc
apploc))

-- | The state is a set of constraints and a counter for generating
-- type names.  This is distinct from the usual counter we use for
-- generating unique names, as these will be user-visible.
data TermTypeState = TermTypeState
  { TermTypeState -> Constraints
stateConstraints :: Constraints,
    TermTypeState -> Int
stateCounter :: !Int,
    -- | Mapping function arguments encountered to
    -- the sizes they ended up generating (when
    -- they could not be substituted directly).
    -- This happens for function arguments that are
    -- not constants or names.
    TermTypeState -> Map SizeSource VName
stateDimTable :: M.Map SizeSource VName,
    TermTypeState -> Map VName NameReason
stateNames :: M.Map VName NameReason,
    TermTypeState -> [Occurrence]
stateOccs :: Occurrences
  }

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

liftTypeM :: TypeM a -> TermTypeM a
liftTypeM :: TypeM a -> TermTypeM a
liftTypeM = ReaderT TermEnv (StateT TermTypeState TypeM) a -> TermTypeM a
forall a.
ReaderT TermEnv (StateT TermTypeState TypeM) a -> TermTypeM a
TermTypeM (ReaderT TermEnv (StateT TermTypeState TypeM) a -> TermTypeM a)
-> (TypeM a -> ReaderT TermEnv (StateT TermTypeState TypeM) a)
-> TypeM a
-> TermTypeM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StateT TermTypeState TypeM a
-> ReaderT TermEnv (StateT TermTypeState TypeM) a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT TermTypeState TypeM a
 -> ReaderT TermEnv (StateT TermTypeState TypeM) a)
-> (TypeM a -> StateT TermTypeState TypeM a)
-> TypeM a
-> ReaderT TermEnv (StateT TermTypeState TypeM) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeM a -> StateT TermTypeState TypeM a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

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

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

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

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

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

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

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

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

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

-- | Create a new type name and insert it (unconstrained) in the
-- substitution map.
instantiateTypeParam ::
  Monoid as =>
  SrcLoc ->
  TypeParam ->
  TermTypeM (VName, Subst (RetTypeBase dim as))
instantiateTypeParam :: SrcLoc
-> TypeParam -> TermTypeM (VName, Subst (RetTypeBase dim as))
instantiateTypeParam SrcLoc
loc TypeParam
tparam = do
  Int
i <- TermTypeM Int
incCounter
  let name :: Name
name = String -> Name
nameFromString ((Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Char -> Bool
isAscii (VName -> String
baseString (TypeParam -> VName
forall vn. TypeParamBase vn -> vn
typeParamName TypeParam
tparam)))
  VName
v <- Name -> TermTypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID (Name -> TermTypeM VName) -> Name -> TermTypeM VName
forall a b. (a -> b) -> a -> b
$ Name -> Int -> Name
mkTypeVarName Name
name Int
i
  case TypeParam
tparam of
    TypeParamType Liftedness
x VName
_ SrcLoc
_ -> do
      VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Liftedness -> Usage -> Constraint
NoConstraint Liftedness
x (Usage -> Constraint) -> Usage -> Constraint
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
mkUsage' SrcLoc
loc
      (VName, Subst (RetTypeBase dim as))
-> TermTypeM (VName, Subst (RetTypeBase dim as))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, [TypeParam] -> RetTypeBase dim as -> Subst (RetTypeBase dim as)
forall t. [TypeParam] -> t -> Subst t
Subst [] (RetTypeBase dim as -> Subst (RetTypeBase dim as))
-> RetTypeBase dim as -> Subst (RetTypeBase dim as)
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase dim as -> RetTypeBase dim as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim as -> RetTypeBase dim as)
-> TypeBase dim as -> RetTypeBase dim as
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar as
forall a. Monoid a => a
mempty Uniqueness
Nonunique (VName -> TypeName
typeName VName
v) [])
    TypeParamDim {} -> do
      VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Maybe (DimDecl VName) -> Usage -> Constraint
Size Maybe (DimDecl VName)
forall a. Maybe a
Nothing (Usage -> Constraint) -> Usage -> Constraint
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
mkUsage' SrcLoc
loc
      (VName, Subst (RetTypeBase dim as))
-> TermTypeM (VName, Subst (RetTypeBase dim as))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, DimDecl VName -> Subst (RetTypeBase dim as)
forall t. DimDecl VName -> Subst t
SizeSubst (DimDecl VName -> Subst (RetTypeBase dim as))
-> DimDecl VName -> Subst (RetTypeBase dim as)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v)

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

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

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

instance MonadTypeChecker TermTypeM where
  warn :: loc -> Doc -> TermTypeM ()
warn loc
loc Doc
problem = TypeM () -> TermTypeM ()
forall a. TypeM a -> TermTypeM a
liftTypeM (TypeM () -> TermTypeM ()) -> TypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ loc -> Doc -> TypeM ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc -> m ()
warn loc
loc Doc
problem
  newName :: VName -> TermTypeM VName
newName = TypeM VName -> TermTypeM VName
forall a. TypeM a -> TermTypeM a
liftTypeM (TypeM VName -> TermTypeM VName)
-> (VName -> TypeM VName) -> VName -> TermTypeM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> TypeM VName
forall (m :: * -> *). MonadTypeChecker m => VName -> m VName
newName
  newID :: Name -> TermTypeM VName
newID = TypeM VName -> TermTypeM VName
forall a. TypeM a -> TermTypeM a
liftTypeM (TypeM VName -> TermTypeM VName)
-> (Name -> TypeM VName) -> Name -> TermTypeM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> TypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newID

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

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

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

  bindVal :: VName -> BoundV -> TermTypeM a -> TermTypeM a
bindVal VName
v (TypeM.BoundV [TypeParam]
tps StructType
t) = (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope ((TermScope -> TermScope) -> TermTypeM a -> TermTypeM a)
-> (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
forall a b. (a -> b) -> a -> b
$ \TermScope
scope ->
    TermScope
scope {scopeVtable :: Map VName ValBinding
scopeVtable = VName -> ValBinding -> Map VName ValBinding -> Map VName ValBinding
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v ValBinding
vb (Map VName ValBinding -> Map VName ValBinding)
-> Map VName ValBinding -> Map VName ValBinding
forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope}
    where
      vb :: ValBinding
vb = Locality -> [TypeParam] -> PatType -> ValBinding
BoundV Locality
Local [TypeParam]
tps (PatType -> ValBinding) -> PatType -> ValBinding
forall a b. (a -> b) -> a -> b
$ StructType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct StructType
t

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

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

  lookupVar :: SrcLoc -> QualName Name -> TermTypeM (QualName VName, PatType)
lookupVar SrcLoc
loc QualName Name
qn = do
    Env
outer_env <- TypeM Env -> TermTypeM Env
forall a. TypeM a -> TermTypeM a
liftTypeM TypeM Env
askEnv
    (TermScope
scope, qn' :: QualName VName
qn'@(QualName [VName]
qs VName
name)) <- Namespace
-> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
checkQualNameWithEnv Namespace
Term QualName Name
qn SrcLoc
loc
    let usage :: Usage
usage = SrcLoc -> String -> Usage
mkUsage SrcLoc
loc (String -> Usage) -> String -> Usage
forall a b. (a -> b) -> a -> b
$ String
"use of " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ShowS
quote (QualName Name -> String
forall a. Pretty a => a -> String
pretty QualName Name
qn)

    PatType
t <- case VName -> Map VName ValBinding -> Maybe ValBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName ValBinding -> Maybe ValBinding)
-> Map VName ValBinding -> Maybe ValBinding
forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope of
      Maybe ValBinding
Nothing ->
        SrcLoc -> Notes -> Doc -> TermTypeM PatType
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> TermTypeM PatType) -> Doc -> TermTypeM PatType
forall a b. (a -> b) -> a -> b
$
          Doc
"Unknown variable" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr QualName Name
qn) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."
      Just (WasConsumed SrcLoc
wloc) -> VName -> SrcLoc -> SrcLoc -> TermTypeM PatType
forall a. VName -> SrcLoc -> SrcLoc -> TermTypeM a
useAfterConsume VName
name SrcLoc
loc SrcLoc
wloc
      Just (BoundV Locality
_ [TypeParam]
tparams PatType
t)
        | String
"_" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` VName -> String
baseString VName
name -> SrcLoc -> QualName Name -> TermTypeM PatType
forall (m :: * -> *) a.
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m a
underscoreUse SrcLoc
loc QualName Name
qn
        | Bool
otherwise -> do
          ([VName]
tnames, PatType
t') <- SrcLoc -> [TypeParam] -> PatType -> TermTypeM ([VName], PatType)
instantiateTypeScheme SrcLoc
loc [TypeParam]
tparams PatType
t
          PatType -> TermTypeM PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> TermTypeM PatType) -> PatType -> TermTypeM PatType
forall a b. (a -> b) -> a -> b
$ Env -> [VName] -> [VName] -> PatType -> PatType
forall as.
Env
-> [VName]
-> [VName]
-> TypeBase (DimDecl VName) as
-> TypeBase (DimDecl VName) as
qualifyTypeVars Env
outer_env [VName]
tnames [VName]
qs PatType
t'
      Just ValBinding
EqualityF -> do
        PatType
argtype <- SrcLoc -> Name -> TermTypeM PatType
forall (m :: * -> *) als dim.
(MonadUnify m, Monoid als) =>
SrcLoc -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
        Usage -> PatType -> TermTypeM ()
forall (m :: * -> *) dim as.
(MonadUnify m, Pretty (ShapeDecl dim), Monoid as) =>
Usage -> TypeBase dim as -> m ()
equalityType Usage
usage PatType
argtype
        PatType -> TermTypeM PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> TermTypeM PatType) -> PatType -> TermTypeM PatType
forall a b. (a -> b) -> a -> b
$
          ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> (PatType -> ScalarTypeBase (DimDecl VName) Aliasing)
-> PatType
-> PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliasing
-> PName
-> PatType
-> RetTypeBase (DimDecl VName) Aliasing
-> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
Unnamed PatType
argtype (RetTypeBase (DimDecl VName) Aliasing
 -> ScalarTypeBase (DimDecl VName) Aliasing)
-> (PatType -> RetTypeBase (DimDecl VName) Aliasing)
-> PatType
-> ScalarTypeBase (DimDecl VName) Aliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (PatType -> PatType) -> PatType -> PatType
forall a b. (a -> b) -> a -> b
$
            ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Aliasing
-> PName
-> PatType
-> RetTypeBase (DimDecl VName) Aliasing
-> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
Unnamed PatType
argtype (RetTypeBase (DimDecl VName) Aliasing
 -> ScalarTypeBase (DimDecl VName) Aliasing)
-> RetTypeBase (DimDecl VName) Aliasing
-> ScalarTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ [VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (PatType -> RetTypeBase (DimDecl VName) Aliasing)
-> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
Bool
      Just (OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt) -> do
        StructType
argtype <- SrcLoc -> Name -> TermTypeM StructType
forall (m :: * -> *) als dim.
(MonadUnify m, Monoid als) =>
SrcLoc -> Name -> m (TypeBase dim als)
newTypeVar SrcLoc
loc Name
"t"
        [PrimType] -> Usage -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
[PrimType] -> Usage -> StructType -> m ()
mustBeOneOf [PrimType]
ts Usage
usage StructType
argtype
        let ([StructType]
pts', StructType
rt') = StructType
-> [Maybe PrimType] -> Maybe PrimType -> ([StructType], StructType)
forall dim as.
TypeBase dim as
-> [Maybe PrimType]
-> Maybe PrimType
-> ([TypeBase dim ()], TypeBase dim ())
instOverloaded StructType
argtype [Maybe PrimType]
pts Maybe PrimType
rt
            arrow :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
arrow TypeBase dim as
xt TypeBase dim as
yt = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty PName
Unnamed TypeBase dim as
xt (RetTypeBase dim as -> ScalarTypeBase dim as)
-> RetTypeBase dim as -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase dim as -> RetTypeBase dim as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase dim as
yt
        PatType -> TermTypeM PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> TermTypeM PatType) -> PatType -> TermTypeM PatType
forall a b. (a -> b) -> a -> b
$ StructType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct (StructType -> PatType) -> StructType -> PatType
forall a b. (a -> b) -> a -> b
$ (StructType -> StructType -> StructType)
-> StructType -> [StructType] -> StructType
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr StructType -> StructType -> StructType
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
arrow StructType
rt' [StructType]
pts'

    Ident -> TermTypeM ()
observe (Ident -> TermTypeM ()) -> Ident -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ VName -> Info PatType -> SrcLoc -> Ident
forall (f :: * -> *) vn.
vn -> f PatType -> SrcLoc -> IdentBase f vn
Ident VName
name (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t) SrcLoc
loc
    (QualName VName, PatType) -> TermTypeM (QualName VName, PatType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName
qn', PatType
t)
    where
      instOverloaded :: TypeBase dim as
-> [Maybe PrimType]
-> Maybe PrimType
-> ([TypeBase dim ()], TypeBase dim ())
instOverloaded TypeBase dim as
argtype [Maybe PrimType]
pts Maybe PrimType
rt =
        ( (Maybe PrimType -> TypeBase dim ())
-> [Maybe PrimType] -> [TypeBase dim ()]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase dim ()
-> (PrimType -> TypeBase dim ())
-> Maybe PrimType
-> TypeBase dim ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TypeBase dim as -> TypeBase dim ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase dim as
argtype) (ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> (PrimType -> ScalarTypeBase dim ())
-> PrimType
-> TypeBase dim ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> ScalarTypeBase dim ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim)) [Maybe PrimType]
pts,
          TypeBase dim ()
-> (PrimType -> TypeBase dim ())
-> Maybe PrimType
-> TypeBase dim ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TypeBase dim as -> TypeBase dim ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase dim as
argtype) (ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> (PrimType -> ScalarTypeBase dim ())
-> PrimType
-> TypeBase dim ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> ScalarTypeBase dim ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim) Maybe PrimType
rt
        )

  checkNamedDim :: SrcLoc -> QualName Name -> TermTypeM (QualName VName)
checkNamedDim SrcLoc
loc QualName Name
v = do
    (QualName VName
v', PatType
t) <- SrcLoc -> QualName Name -> TermTypeM (QualName VName, PatType)
forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m (QualName VName, PatType)
lookupVar SrcLoc
loc QualName Name
v
    Checking -> TermTypeM () -> TermTypeM ()
forall a. Checking -> TermTypeM a -> TermTypeM a
onFailure ([StructType] -> StructType -> Checking
CheckingRequired [ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> StructType)
-> ScalarTypeBase (DimDecl VName) () -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) ())
-> PrimType -> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64] (PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)) (TermTypeM () -> TermTypeM ()) -> TermTypeM () -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
      Usage -> StructType -> StructType -> TermTypeM ()
forall (m :: * -> *).
MonadUnify m =>
Usage -> StructType -> StructType -> m ()
unify (SrcLoc -> String -> Usage
mkUsage SrcLoc
loc String
"use as array size") (PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) (StructType -> TermTypeM ()) -> StructType -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$
        ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> StructType)
-> ScalarTypeBase (DimDecl VName) () -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) ())
-> PrimType -> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    QualName VName -> TermTypeM (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure QualName VName
v'

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

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

extSize :: SrcLoc -> SizeSource -> TermTypeM (DimDecl VName, Maybe VName)
extSize :: SrcLoc -> SizeSource -> TermTypeM (DimDecl VName, Maybe VName)
extSize SrcLoc
loc SizeSource
e = do
  Maybe VName
prev <- (TermTypeState -> Maybe VName) -> TermTypeM (Maybe VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((TermTypeState -> Maybe VName) -> TermTypeM (Maybe VName))
-> (TermTypeState -> Maybe VName) -> TermTypeM (Maybe VName)
forall a b. (a -> b) -> a -> b
$ SizeSource -> Map SizeSource VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup SizeSource
e (Map SizeSource VName -> Maybe VName)
-> (TermTypeState -> Map SizeSource VName)
-> TermTypeState
-> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermTypeState -> Map SizeSource VName
stateDimTable
  case Maybe VName
prev of
    Maybe VName
Nothing -> do
      let rsrc :: RigidSource
rsrc = case SizeSource
e of
            SourceArg (FName Maybe (QualName VName)
fname) ExpBase NoInfo VName
e' ->
              Maybe (QualName VName) -> String -> RigidSource
RigidArg Maybe (QualName VName)
fname (String -> RigidSource) -> String -> RigidSource
forall a b. (a -> b) -> a -> b
$ ExpBase NoInfo VName -> String
forall a. Pretty a => a -> String
prettyOneLine ExpBase NoInfo VName
e'
            SourceBound ExpBase NoInfo VName
e' ->
              String -> RigidSource
RigidBound (String -> RigidSource) -> String -> RigidSource
forall a b. (a -> b) -> a -> b
$ ExpBase NoInfo VName -> String
forall a. Pretty a => a -> String
prettyOneLine ExpBase NoInfo VName
e'
            SourceSlice Maybe (DimDecl VName)
d Maybe (ExpBase NoInfo VName)
i Maybe (ExpBase NoInfo VName)
j Maybe (ExpBase NoInfo VName)
s ->
              Maybe (DimDecl VName) -> String -> RigidSource
RigidSlice Maybe (DimDecl VName)
d (String -> RigidSource) -> String -> RigidSource
forall a b. (a -> b) -> a -> b
$ DimIndexBase NoInfo VName -> String
forall a. Pretty a => a -> String
prettyOneLine (DimIndexBase NoInfo VName -> String)
-> DimIndexBase NoInfo VName -> String
forall a b. (a -> b) -> a -> b
$ Maybe (ExpBase NoInfo VName)
-> Maybe (ExpBase NoInfo VName)
-> Maybe (ExpBase NoInfo VName)
-> DimIndexBase NoInfo VName
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice Maybe (ExpBase NoInfo VName)
i Maybe (ExpBase NoInfo VName)
j Maybe (ExpBase NoInfo VName)
s
      VName
d <- SrcLoc -> Rigidity -> Name -> TermTypeM VName
forall (m :: * -> *).
MonadUnify m =>
SrcLoc -> Rigidity -> Name -> m VName
newDimVar SrcLoc
loc (RigidSource -> Rigidity
Rigid RigidSource
rsrc) Name
"n"
      (TermTypeState -> TermTypeState) -> TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TermTypeState -> TermTypeState) -> TermTypeM ())
-> (TermTypeState -> TermTypeState) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateDimTable :: Map SizeSource VName
stateDimTable = SizeSource -> VName -> Map SizeSource VName -> Map SizeSource VName
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert SizeSource
e VName
d (Map SizeSource VName -> Map SizeSource VName)
-> Map SizeSource VName -> Map SizeSource VName
forall a b. (a -> b) -> a -> b
$ TermTypeState -> Map SizeSource VName
stateDimTable TermTypeState
s}
      (DimDecl VName, Maybe VName)
-> TermTypeM (DimDecl VName, Maybe VName)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d,
          VName -> Maybe VName
forall a. a -> Maybe a
Just VName
d
        )
    Just VName
d ->
      (DimDecl VName, Maybe VName)
-> TermTypeM (DimDecl VName, Maybe VName)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d,
          VName -> Maybe VName
forall a. a -> Maybe a
Just VName
d
        )

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

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

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

newArrayType :: SrcLoc -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType :: SrcLoc -> Name -> Int -> TermTypeM (StructType, StructType)
newArrayType SrcLoc
loc Name
desc Int
r = do
  VName
v <- Name -> TermTypeM VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
desc
  VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Liftedness -> Usage -> Constraint
NoConstraint Liftedness
Unlifted (Usage -> Constraint) -> Usage -> Constraint
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
mkUsage' SrcLoc
loc
  [VName]
dims <- Int -> TermTypeM VName -> TermTypeM [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
r (TermTypeM VName -> TermTypeM [VName])
-> TermTypeM VName -> TermTypeM [VName]
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Rigidity -> Name -> TermTypeM VName
forall (m :: * -> *).
MonadUnify m =>
SrcLoc -> Rigidity -> Name -> m VName
newDimVar SrcLoc
loc Rigidity
Nonrigid Name
"dim"
  let rowt :: ScalarTypeBase dim ()
rowt = ()
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim ()
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> TypeName
typeName VName
v) []
  (StructType, StructType) -> TermTypeM (StructType, StructType)
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> StructType
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
rowt ([DimDecl VName] -> ShapeDecl (DimDecl VName)
forall dim. [dim] -> ShapeDecl dim
ShapeDecl ([DimDecl VName] -> ShapeDecl (DimDecl VName))
-> [DimDecl VName] -> ShapeDecl (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ (VName -> DimDecl VName) -> [VName] -> [DimDecl VName]
forall a b. (a -> b) -> [a] -> [b]
map (QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> (VName -> QualName VName) -> VName -> DimDecl VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName) [VName]
dims),
      ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
rowt
    )

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

-- | Replace all type variables with their concrete types.
updateTypes :: ASTMappable e => e -> TermTypeM e
updateTypes :: e -> TermTypeM e
updateTypes = ASTMapper TermTypeM -> e -> TermTypeM e
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper TermTypeM
tv
  where
    tv :: ASTMapper TermTypeM
tv =
      ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (StructType -> m StructType)
-> (PatType -> m PatType)
-> (RetTypeBase (DimDecl VName) ()
    -> m (RetTypeBase (DimDecl VName) ()))
-> (RetTypeBase (DimDecl VName) Aliasing
    -> m (RetTypeBase (DimDecl VName) Aliasing))
-> ASTMapper m
ASTMapper
        { mapOnExp :: Exp -> TermTypeM Exp
mapOnExp = ASTMapper TermTypeM -> Exp -> TermTypeM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper TermTypeM
tv,
          mapOnName :: VName -> TermTypeM VName
mapOnName = VName -> TermTypeM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnQualName :: QualName VName -> TermTypeM (QualName VName)
mapOnQualName = QualName VName -> TermTypeM (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnStructType :: StructType -> TermTypeM StructType
mapOnStructType = StructType -> TermTypeM StructType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully,
          mapOnPatType :: PatType -> TermTypeM PatType
mapOnPatType = PatType -> TermTypeM PatType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully,
          mapOnStructRetType :: RetTypeBase (DimDecl VName) ()
-> TermTypeM (RetTypeBase (DimDecl VName) ())
mapOnStructRetType = RetTypeBase (DimDecl VName) ()
-> TermTypeM (RetTypeBase (DimDecl VName) ())
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully,
          mapOnPatRetType :: RetTypeBase (DimDecl VName) Aliasing
-> TermTypeM (RetTypeBase (DimDecl VName) Aliasing)
mapOnPatRetType = RetTypeBase (DimDecl VName) Aliasing
-> TermTypeM (RetTypeBase (DimDecl VName) Aliasing)
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully
        }

--- Basic checking

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

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

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

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

  -- Observe the sizes so we do not get any warnings about them not
  -- being used.
  (DimDecl VName -> TermTypeM ()) -> [DimDecl VName] -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ DimDecl VName -> TermTypeM ()
observeDim ([DimDecl VName] -> TermTypeM ())
-> [DimDecl VName] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ StructType -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims StructType
st
  (TypeExp VName, [VName], RetTypeBase (DimDecl VName) ())
-> TermTypeM
     (TypeExp VName, [VName], RetTypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp VName
te', [VName]
svars, [VName] -> StructType -> RetTypeBase (DimDecl VName) ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims StructType
st)
  where
    observeDim :: DimDecl VName -> TermTypeM ()
observeDim (NamedDim QualName VName
v) =
      Ident -> TermTypeM ()
observe (Ident -> TermTypeM ()) -> Ident -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ VName -> Info PatType -> SrcLoc -> Ident
forall (f :: * -> *) vn.
vn -> f PatType -> SrcLoc -> IdentBase f vn
Ident (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) Aliasing)
-> PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64) SrcLoc
forall a. Monoid a => a
mempty
    observeDim DimDecl VName
_ = () -> TermTypeM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

checkTypeExpNonrigid :: TypeExp Name -> TermTypeM (TypeExp VName, StructType, [VName])
checkTypeExpNonrigid :: TypeExp Name -> TermTypeM (TypeExp VName, StructType, [VName])
checkTypeExpNonrigid TypeExp Name
te = do
  (TypeExp VName
te', [VName]
svars, RetType [VName]
dims StructType
st) <- TypeExp Name
-> TermTypeM
     (TypeExp VName, [VName], RetTypeBase (DimDecl VName) ())
termCheckTypeExp TypeExp Name
te
  [VName] -> (VName -> TermTypeM ()) -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName]
svars [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims) ((VName -> TermTypeM ()) -> TermTypeM ())
-> (VName -> TermTypeM ()) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \VName
v ->
    VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ Maybe (DimDecl VName) -> Usage -> Constraint
Size Maybe (DimDecl VName)
forall a. Maybe a
Nothing (Usage -> Constraint) -> Usage -> Constraint
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Usage
mkUsage' (SrcLoc -> Usage) -> SrcLoc -> Usage
forall a b. (a -> b) -> a -> b
$ TypeExp Name -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeExp Name
te
  (TypeExp VName, StructType, [VName])
-> TermTypeM (TypeExp VName, StructType, [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp VName
te', StructType
st, [VName]
svars [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims)

checkTypeExpRigid ::
  TypeExp Name ->
  RigidSource ->
  TermTypeM (TypeExp VName, StructType, [VName])
checkTypeExpRigid :: TypeExp Name
-> RigidSource -> TermTypeM (TypeExp VName, StructType, [VName])
checkTypeExpRigid TypeExp Name
te RigidSource
rsrc = do
  (TypeExp VName
te', [VName]
svars, RetType [VName]
dims StructType
st) <- TypeExp Name
-> TermTypeM
     (TypeExp VName, [VName], RetTypeBase (DimDecl VName) ())
termCheckTypeExp TypeExp Name
te
  [VName] -> (VName -> TermTypeM ()) -> TermTypeM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName]
svars [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims) ((VName -> TermTypeM ()) -> TermTypeM ())
-> (VName -> TermTypeM ()) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \VName
v ->
    VName -> Constraint -> TermTypeM ()
constrain VName
v (Constraint -> TermTypeM ()) -> Constraint -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ SrcLoc -> RigidSource -> Constraint
UnknowableSize (TypeExp Name -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeExp Name
te) RigidSource
rsrc
  (TypeExp VName, StructType, [VName])
-> TermTypeM (TypeExp VName, StructType, [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp VName
te', StructType
st, [VName]
svars [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims)

--- Sizes

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

maybeDimFromExp :: Exp -> Maybe (DimDecl VName)
maybeDimFromExp :: Exp -> Maybe (DimDecl VName)
maybeDimFromExp (Var QualName VName
v Info PatType
_ SrcLoc
_) = DimDecl VName -> Maybe (DimDecl VName)
forall a. a -> Maybe a
Just (DimDecl VName -> Maybe (DimDecl VName))
-> DimDecl VName -> Maybe (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v
maybeDimFromExp (Parens Exp
e SrcLoc
_) = Exp -> Maybe (DimDecl VName)
maybeDimFromExp Exp
e
maybeDimFromExp (QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) = Exp -> Maybe (DimDecl VName)
maybeDimFromExp Exp
e
maybeDimFromExp Exp
e = Int -> DimDecl VName
forall vn. Int -> DimDecl vn
ConstDim (Int -> DimDecl VName) -> (Int64 -> Int) -> Int64 -> DimDecl VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> DimDecl VName) -> Maybe Int64 -> Maybe (DimDecl VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Maybe Int64
isInt64 Exp
e

dimFromExp :: (Exp -> SizeSource) -> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromExp :: (Exp -> SizeSource)
-> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromExp Exp -> SizeSource
rf (Parens Exp
e SrcLoc
_) = (Exp -> SizeSource)
-> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromExp Exp -> SizeSource
rf Exp
e
dimFromExp Exp -> SizeSource
rf (QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) = (Exp -> SizeSource)
-> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromExp Exp -> SizeSource
rf Exp
e
dimFromExp Exp -> SizeSource
rf Exp
e
  | Just DimDecl VName
d <- Exp -> Maybe (DimDecl VName)
maybeDimFromExp Exp
e =
    (DimDecl VName, Maybe VName)
-> TermTypeM (DimDecl VName, Maybe VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName
d, Maybe VName
forall a. Maybe a
Nothing)
  | Bool
otherwise =
    SrcLoc -> SizeSource -> TermTypeM (DimDecl VName, Maybe VName)
extSize (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e) (SizeSource -> TermTypeM (DimDecl VName, Maybe VName))
-> SizeSource -> TermTypeM (DimDecl VName, Maybe VName)
forall a b. (a -> b) -> a -> b
$ Exp -> SizeSource
rf Exp
e

dimFromArg :: Maybe (QualName VName) -> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromArg :: Maybe (QualName VName)
-> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromArg Maybe (QualName VName)
fname = (Exp -> SizeSource)
-> Exp -> TermTypeM (DimDecl VName, Maybe VName)
dimFromExp ((Exp -> SizeSource)
 -> Exp -> TermTypeM (DimDecl VName, Maybe VName))
-> (Exp -> SizeSource)
-> Exp
-> TermTypeM (DimDecl VName, Maybe VName)
forall a b. (a -> b) -> a -> b
$ FName -> ExpBase NoInfo VName -> SizeSource
SourceArg (Maybe (QualName VName) -> FName
FName Maybe (QualName VName)
fname) (ExpBase NoInfo VName -> SizeSource)
-> (Exp -> ExpBase NoInfo VName) -> Exp -> SizeSource
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> ExpBase NoInfo VName
bareExp

-- | Any argument sizes created with 'extSize' inside the given action
-- will be removed once the action finishes.  This is to ensure that
-- just because e.g. @n+1@ appears as a size in one branch of a
-- conditional, that doesn't mean it's also available in the other
-- branch.
noSizeEscape :: TermTypeM a -> TermTypeM a
noSizeEscape :: TermTypeM a -> TermTypeM a
noSizeEscape TermTypeM a
m = do
  Map SizeSource VName
dimtable <- (TermTypeState -> Map SizeSource VName)
-> TermTypeM (Map SizeSource VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TermTypeState -> Map SizeSource VName
stateDimTable
  a
x <- TermTypeM a
m
  (TermTypeState -> TermTypeState) -> TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TermTypeState -> TermTypeState) -> TermTypeM ())
-> (TermTypeState -> TermTypeState) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateDimTable :: Map SizeSource VName
stateDimTable = Map SizeSource VName
dimtable}
  a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

--- Control flow

tapOccurrences :: TermTypeM a -> TermTypeM (a, Occurrences)
tapOccurrences :: TermTypeM a -> TermTypeM (a, [Occurrence])
tapOccurrences TermTypeM a
m = do
  (a
x, [Occurrence]
occs) <- TermTypeM a -> TermTypeM (a, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences TermTypeM a
m
  [Occurrence] -> TermTypeM ()
occur [Occurrence]
occs
  (a, [Occurrence]) -> TermTypeM (a, [Occurrence])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, [Occurrence]
occs)

collectOccurrences :: TermTypeM a -> TermTypeM (a, Occurrences)
collectOccurrences :: TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences TermTypeM a
m = do
  [Occurrence]
old <- (TermTypeState -> [Occurrence]) -> TermTypeM [Occurrence]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TermTypeState -> [Occurrence]
stateOccs
  (TermTypeState -> TermTypeState) -> TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TermTypeState -> TermTypeState) -> TermTypeM ())
-> (TermTypeState -> TermTypeState) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateOccs :: [Occurrence]
stateOccs = [Occurrence]
forall a. Monoid a => a
mempty}
  a
x <- TermTypeM a
m
  [Occurrence]
new <- (TermTypeState -> [Occurrence]) -> TermTypeM [Occurrence]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TermTypeState -> [Occurrence]
stateOccs
  (TermTypeState -> TermTypeState) -> TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TermTypeState -> TermTypeState) -> TermTypeM ())
-> (TermTypeState -> TermTypeState) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateOccs :: [Occurrence]
stateOccs = [Occurrence]
old}
  (a, [Occurrence]) -> TermTypeM (a, [Occurrence])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, [Occurrence]
new)

alternative :: TermTypeM a -> TermTypeM b -> TermTypeM (a, b)
alternative :: TermTypeM a -> TermTypeM b -> TermTypeM (a, b)
alternative TermTypeM a
m1 TermTypeM b
m2 = do
  (a
x, [Occurrence]
occurs1) <- TermTypeM a -> TermTypeM (a, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences (TermTypeM a -> TermTypeM (a, [Occurrence]))
-> TermTypeM a -> TermTypeM (a, [Occurrence])
forall a b. (a -> b) -> a -> b
$ TermTypeM a -> TermTypeM a
forall a. TermTypeM a -> TermTypeM a
noSizeEscape TermTypeM a
m1
  (b
y, [Occurrence]
occurs2) <- TermTypeM b -> TermTypeM (b, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences (TermTypeM b -> TermTypeM (b, [Occurrence]))
-> TermTypeM b -> TermTypeM (b, [Occurrence])
forall a b. (a -> b) -> a -> b
$ TermTypeM b -> TermTypeM b
forall a. TermTypeM a -> TermTypeM a
noSizeEscape TermTypeM b
m2
  [Occurrence] -> TermTypeM ()
checkOccurrences [Occurrence]
occurs1
  [Occurrence] -> TermTypeM ()
checkOccurrences [Occurrence]
occurs2
  [Occurrence] -> TermTypeM ()
occur ([Occurrence] -> TermTypeM ()) -> [Occurrence] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ [Occurrence]
occurs1 [Occurrence] -> [Occurrence] -> [Occurrence]
`altOccurrences` [Occurrence]
occurs2
  (a, b) -> TermTypeM (a, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, b
y)

sequentially :: TermTypeM a -> (a -> Occurrences -> TermTypeM b) -> TermTypeM b
sequentially :: TermTypeM a -> (a -> [Occurrence] -> TermTypeM b) -> TermTypeM b
sequentially TermTypeM a
m1 a -> [Occurrence] -> TermTypeM b
m2 = do
  (a
a, [Occurrence]
m1flow) <- TermTypeM a -> TermTypeM (a, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences TermTypeM a
m1
  (b
b, [Occurrence]
m2flow) <- TermTypeM b -> TermTypeM (b, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences (TermTypeM b -> TermTypeM (b, [Occurrence]))
-> TermTypeM b -> TermTypeM (b, [Occurrence])
forall a b. (a -> b) -> a -> b
$ a -> [Occurrence] -> TermTypeM b
m2 a
a [Occurrence]
m1flow
  [Occurrence] -> TermTypeM ()
occur ([Occurrence] -> TermTypeM ()) -> [Occurrence] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ [Occurrence]
m1flow [Occurrence] -> [Occurrence] -> [Occurrence]
`seqOccurrences` [Occurrence]
m2flow
  b -> TermTypeM b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b

--- Consumption

occur :: Occurrences -> TermTypeM ()
occur :: [Occurrence] -> TermTypeM ()
occur [Occurrence]
occs = (TermTypeState -> TermTypeState) -> TermTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TermTypeState -> TermTypeState) -> TermTypeM ())
-> (TermTypeState -> TermTypeState) -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ \TermTypeState
s -> TermTypeState
s {stateOccs :: [Occurrence]
stateOccs = TermTypeState -> [Occurrence]
stateOccs TermTypeState
s [Occurrence] -> [Occurrence] -> [Occurrence]
forall a. Semigroup a => a -> a -> a
<> [Occurrence]
occs}

-- | Proclaim that we have made read-only use of the given variable.
observe :: Ident -> TermTypeM ()
observe :: Ident -> TermTypeM ()
observe (Ident VName
nm (Info PatType
t) SrcLoc
loc) =
  let als :: Aliasing
als = VName -> Alias
AliasBound VName
nm Alias -> Aliasing -> Aliasing
forall a. Ord a => a -> Set a -> Set a
`S.insert` PatType -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases PatType
t
   in [Occurrence] -> TermTypeM ()
occur [Aliasing -> SrcLoc -> Occurrence
observation Aliasing
als SrcLoc
loc]

onlySelfAliasing :: TermTypeM a -> TermTypeM a
onlySelfAliasing :: TermTypeM a -> TermTypeM a
onlySelfAliasing = (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope (\TermScope
scope -> TermScope
scope {scopeVtable :: Map VName ValBinding
scopeVtable = (VName -> ValBinding -> ValBinding)
-> Map VName ValBinding -> Map VName ValBinding
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey VName -> ValBinding -> ValBinding
set (Map VName ValBinding -> Map VName ValBinding)
-> Map VName ValBinding -> Map VName ValBinding
forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope})
  where
    set :: VName -> ValBinding -> ValBinding
set VName
k (BoundV Locality
l [TypeParam]
tparams PatType
t) =
      Locality -> [TypeParam] -> PatType -> ValBinding
BoundV Locality
l [TypeParam]
tparams (PatType -> ValBinding) -> PatType -> ValBinding
forall a b. (a -> b) -> a -> b
$
        PatType
t PatType -> (Aliasing -> Aliasing) -> PatType
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` Aliasing -> Aliasing -> Aliasing
forall a. Ord a => Set a -> Set a -> Set a
S.intersection (Alias -> Aliasing
forall a. a -> Set a
S.singleton (VName -> Alias
AliasBound VName
k))
    set VName
_ (OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt) = [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> ValBinding
OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt
    set VName
_ ValBinding
EqualityF = ValBinding
EqualityF
    set VName
_ (WasConsumed SrcLoc
loc) = SrcLoc -> ValBinding
WasConsumed SrcLoc
loc

-- | Enter a context where nothing outside can be consumed (i.e. the
-- body of a function definition).
noUnique :: TermTypeM a -> TermTypeM a
noUnique :: TermTypeM a -> TermTypeM a
noUnique TermTypeM a
m = do
  (a
x, [Occurrence]
occs) <- TermTypeM a -> TermTypeM (a, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences (TermTypeM a -> TermTypeM (a, [Occurrence]))
-> TermTypeM a -> TermTypeM (a, [Occurrence])
forall a b. (a -> b) -> a -> b
$ (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope TermScope -> TermScope
f TermTypeM a
m
  [Occurrence] -> TermTypeM ()
checkOccurrences [Occurrence]
occs
  [Occurrence] -> TermTypeM ()
occur ([Occurrence] -> TermTypeM ()) -> [Occurrence] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ ([Occurrence], [Occurrence]) -> [Occurrence]
forall a b. (a, b) -> a
fst (([Occurrence], [Occurrence]) -> [Occurrence])
-> ([Occurrence], [Occurrence]) -> [Occurrence]
forall a b. (a -> b) -> a -> b
$ [Occurrence] -> ([Occurrence], [Occurrence])
split [Occurrence]
occs
  a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
  where
    f :: TermScope -> TermScope
f TermScope
scope = TermScope
scope {scopeVtable :: Map VName ValBinding
scopeVtable = (ValBinding -> ValBinding)
-> Map VName ValBinding -> Map VName ValBinding
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ValBinding -> ValBinding
set (Map VName ValBinding -> Map VName ValBinding)
-> Map VName ValBinding -> Map VName ValBinding
forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope}

    set :: ValBinding -> ValBinding
set (BoundV Locality
l [TypeParam]
tparams PatType
t) = Locality -> [TypeParam] -> PatType -> ValBinding
BoundV Locality
l [TypeParam]
tparams (PatType -> ValBinding) -> PatType -> ValBinding
forall a b. (a -> b) -> a -> b
$ PatType
t PatType -> Uniqueness -> PatType
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
    set (OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt) = [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> ValBinding
OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rt
    set ValBinding
EqualityF = ValBinding
EqualityF
    set (WasConsumed SrcLoc
loc) = SrcLoc -> ValBinding
WasConsumed SrcLoc
loc

    split :: [Occurrence] -> ([Occurrence], [Occurrence])
split = [(Occurrence, Occurrence)] -> ([Occurrence], [Occurrence])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Occurrence, Occurrence)] -> ([Occurrence], [Occurrence]))
-> ([Occurrence] -> [(Occurrence, Occurrence)])
-> [Occurrence]
-> ([Occurrence], [Occurrence])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Occurrence -> (Occurrence, Occurrence))
-> [Occurrence] -> [(Occurrence, Occurrence)]
forall a b. (a -> b) -> [a] -> [b]
map (\Occurrence
occ -> (Occurrence
occ {consumed :: Maybe Names
consumed = Maybe Names
forall a. Monoid a => a
mempty}, Occurrence
occ {observed :: Names
observed = Names
forall a. Monoid a => a
mempty}))

removeSeminullOccurrences :: TermTypeM a -> TermTypeM a
removeSeminullOccurrences :: TermTypeM a -> TermTypeM a
removeSeminullOccurrences TermTypeM a
m = do
  (a
x, [Occurrence]
occs) <- TermTypeM a -> TermTypeM (a, [Occurrence])
forall a. TermTypeM a -> TermTypeM (a, [Occurrence])
collectOccurrences TermTypeM a
m
  [Occurrence] -> TermTypeM ()
occur ([Occurrence] -> TermTypeM ()) -> [Occurrence] -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ (Occurrence -> Bool) -> [Occurrence] -> [Occurrence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurrence -> Bool) -> Occurrence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurrence -> Bool
seminullOccurrence) [Occurrence]
occs
  a -> TermTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

checkIfConsumable :: SrcLoc -> Aliasing -> TermTypeM ()
checkIfConsumable :: SrcLoc -> Aliasing -> TermTypeM ()
checkIfConsumable SrcLoc
loc Aliasing
als = do
  Map VName ValBinding
vtable <- (TermEnv -> Map VName ValBinding)
-> TermTypeM (Map VName ValBinding)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((TermEnv -> Map VName ValBinding)
 -> TermTypeM (Map VName ValBinding))
-> (TermEnv -> Map VName ValBinding)
-> TermTypeM (Map VName ValBinding)
forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable (TermScope -> Map VName ValBinding)
-> (TermEnv -> TermScope) -> TermEnv -> Map VName ValBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermEnv -> TermScope
termScope
  let consumable :: VName -> Bool
consumable VName
v = case VName -> Map VName ValBinding -> Maybe ValBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName ValBinding
vtable of
        Just (BoundV Locality
Local [TypeParam]
_ PatType
t)
          | PatType -> Int
forall dim as. TypeBase dim as -> Int
arrayRank PatType
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> PatType -> Bool
forall shape as. TypeBase shape as -> Bool
unique PatType
t
          | Scalar TypeVar {} <- PatType
t -> PatType -> Bool
forall shape as. TypeBase shape as -> Bool
unique PatType
t
          | Scalar Arrow {} <- PatType
t -> Bool
False
          | Bool
otherwise -> Bool
True
        Just (BoundV Locality
Global [TypeParam]
_ PatType
_) -> Bool
False
        Maybe ValBinding
_ -> Bool
True
  -- The sort ensures that AliasBound vars are shown before AliasFree.
  case (Alias -> VName) -> [Alias] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Alias -> VName
aliasVar ([Alias] -> [VName]) -> [Alias] -> [VName]
forall a b. (a -> b) -> a -> b
$ [Alias] -> [Alias]
forall a. Ord a => [a] -> [a]
sort ([Alias] -> [Alias]) -> [Alias] -> [Alias]
forall a b. (a -> b) -> a -> b
$ (Alias -> Bool) -> [Alias] -> [Alias]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Alias -> Bool) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
consumable (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) ([Alias] -> [Alias]) -> [Alias] -> [Alias]
forall a b. (a -> b) -> a -> b
$ Aliasing -> [Alias]
forall a. Set a -> [a]
S.toList Aliasing
als of
    VName
v : [VName]
_ -> SrcLoc -> Doc -> TermTypeM ()
forall (m :: * -> *) b. MonadTypeChecker m => SrcLoc -> Doc -> m b
notConsumable SrcLoc
loc (Doc -> TermTypeM ()) -> TermTypeM Doc -> TermTypeM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SrcLoc -> VName -> TermTypeM Doc
describeVar SrcLoc
loc VName
v
    [] -> () -> TermTypeM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Proclaim that we have written to the given variable.
consume :: SrcLoc -> Aliasing -> TermTypeM ()
consume :: SrcLoc -> Aliasing -> TermTypeM ()
consume SrcLoc
loc Aliasing
als = do
  SrcLoc -> Aliasing -> TermTypeM ()
checkIfConsumable SrcLoc
loc Aliasing
als
  [Occurrence] -> TermTypeM ()
occur [Aliasing -> SrcLoc -> Occurrence
consumption Aliasing
als SrcLoc
loc]

-- | Proclaim that we have written to the given variable, and mark
-- accesses to it and all of its aliases as invalid inside the given
-- computation.
consuming :: Ident -> TermTypeM a -> TermTypeM a
consuming :: Ident -> TermTypeM a -> TermTypeM a
consuming (Ident VName
name (Info PatType
t) SrcLoc
loc) TermTypeM a
m = do
  PatType
t' <- PatType -> TermTypeM PatType
forall a (m :: * -> *). (Substitutable a, MonadUnify m) => a -> m a
normTypeFully PatType
t
  SrcLoc -> Aliasing -> TermTypeM ()
consume SrcLoc
loc (Aliasing -> TermTypeM ()) -> Aliasing -> TermTypeM ()
forall a b. (a -> b) -> a -> b
$ VName -> Alias
AliasBound VName
name Alias -> Aliasing -> Aliasing
forall a. Ord a => a -> Set a -> Set a
`S.insert` PatType -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases PatType
t'
  (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
forall a. (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a
localScope TermScope -> TermScope
consume' TermTypeM a
m
  where
    consume' :: TermScope -> TermScope
consume' TermScope
scope =
      TermScope
scope {scopeVtable :: Map VName ValBinding
scopeVtable = VName -> ValBinding -> Map VName ValBinding -> Map VName ValBinding
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name (SrcLoc -> ValBinding
WasConsumed SrcLoc
loc) (Map VName ValBinding -> Map VName ValBinding)
-> Map VName ValBinding -> Map VName ValBinding
forall a b. (a -> b) -> a -> b
$ TermScope -> Map VName ValBinding
scopeVtable TermScope
scope}

-- Running

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

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

    addIntrinsicF :: (a, Intrinsic) -> Maybe (a, ValBinding)
addIntrinsicF (a
name, IntrinsicMonoFun [PrimType]
pts PrimType
t) =
      (a, ValBinding) -> Maybe (a, ValBinding)
forall a. a -> Maybe a
Just (a
name, Locality -> [TypeParam] -> PatType -> ValBinding
BoundV Locality
Global [] (PatType -> ValBinding) -> PatType -> ValBinding
forall a b. (a -> b) -> a -> b
$ PatType -> RetTypeBase (DimDecl VName) Aliasing -> PatType
forall as dim.
Monoid as =>
TypeBase dim as -> RetTypeBase dim as -> TypeBase dim as
arrow PatType
forall dim as. TypeBase dim as
pts' (RetTypeBase (DimDecl VName) Aliasing -> PatType)
-> RetTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ [VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (PatType -> RetTypeBase (DimDecl VName) Aliasing)
-> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ PrimType -> PatType
forall dim as. PrimType -> TypeBase dim as
prim PrimType
t)
      where
        pts' :: TypeBase dim as
pts' = case [PrimType]
pts of
          [PrimType
pt] -> PrimType -> TypeBase dim as
forall dim as. PrimType -> TypeBase dim as
prim PrimType
pt
          [PrimType]
_ -> ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ [TypeBase dim as] -> ScalarTypeBase dim as
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord ([TypeBase dim as] -> ScalarTypeBase dim as)
-> [TypeBase dim as] -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ (PrimType -> TypeBase dim as) -> [PrimType] -> [TypeBase dim as]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> TypeBase dim as
forall dim as. PrimType -> TypeBase dim as
prim [PrimType]
pts
    addIntrinsicF (a
name, IntrinsicOverloadedFun [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rts) =
      (a, ValBinding) -> Maybe (a, ValBinding)
forall a. a -> Maybe a
Just (a
name, [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> ValBinding
OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rts)
    addIntrinsicF (a
name, IntrinsicPolyFun [TypeParam]
tvs [StructType]
pts RetTypeBase (DimDecl VName) ()
rt) =
      (a, ValBinding) -> Maybe (a, ValBinding)
forall a. a -> Maybe a
Just
        ( a
name,
          Locality -> [TypeParam] -> PatType -> ValBinding
BoundV Locality
Global [TypeParam]
tvs (PatType -> ValBinding) -> PatType -> ValBinding
forall a b. (a -> b) -> a -> b
$
            StructType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct (StructType -> PatType) -> StructType -> PatType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> StructType)
-> ScalarTypeBase (DimDecl VName) () -> StructType
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> StructType
-> RetTypeBase (DimDecl VName) ()
-> ScalarTypeBase (DimDecl VName) ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow ()
forall a. Monoid a => a
mempty PName
Unnamed StructType
pts' RetTypeBase (DimDecl VName) ()
rt
        )
      where
        pts' :: StructType
pts' = case [StructType]
pts of
          [StructType
pt] -> StructType
pt
          [StructType]
_ -> ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> StructType)
-> ScalarTypeBase (DimDecl VName) () -> StructType
forall a b. (a -> b) -> a -> b
$ [StructType] -> ScalarTypeBase (DimDecl VName) ()
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord [StructType]
pts
    addIntrinsicF (a
name, Intrinsic
IntrinsicEquality) =
      (a, ValBinding) -> Maybe (a, ValBinding)
forall a. a -> Maybe a
Just (a
name, ValBinding
EqualityF)
    addIntrinsicF (a, Intrinsic)
_ = Maybe (a, ValBinding)
forall a. Maybe a
Nothing

runTermTypeM :: TermTypeM a -> TypeM (a, Occurrences)
runTermTypeM :: TermTypeM a -> TypeM (a, [Occurrence])
runTermTypeM (TermTypeM ReaderT TermEnv (StateT TermTypeState TypeM) a
m) = do
  TermScope
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
envToTermScope (Env -> TermScope) -> TypeM Env -> TypeM TermScope
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeM Env
askEnv
  let initial_tenv :: TermEnv
initial_tenv =
        TermEnv :: TermScope -> Maybe Checking -> Int -> TermEnv
TermEnv
          { termScope :: TermScope
termScope = TermScope
initial_scope,
            termChecking :: Maybe Checking
termChecking = Maybe Checking
forall a. Maybe a
Nothing,
            termLevel :: Int
termLevel = Int
0
          }
  (TermTypeState -> [Occurrence])
-> (a, TermTypeState) -> (a, [Occurrence])
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second TermTypeState -> [Occurrence]
stateOccs
    ((a, TermTypeState) -> (a, [Occurrence]))
-> TypeM (a, TermTypeState) -> TypeM (a, [Occurrence])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT TermTypeState TypeM a
-> TermTypeState -> TypeM (a, TermTypeState)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT
      (ReaderT TermEnv (StateT TermTypeState TypeM) a
-> TermEnv -> StateT TermTypeState TypeM a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT TermEnv (StateT TermTypeState TypeM) a
m TermEnv
initial_tenv)
      (Constraints
-> Int
-> Map SizeSource VName
-> Map VName NameReason
-> [Occurrence]
-> TermTypeState
TermTypeState Constraints
forall a. Monoid a => a
mempty Int
0 Map SizeSource VName
forall a. Monoid a => a
mempty Map VName NameReason
forall a. Monoid a => a
mempty [Occurrence]
forall a. Monoid a => a
mempty)