{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE TypeFamilies #-}

-- | The type checker checks whether the program is type-consistent.
module Futhark.IR.TypeCheck
  ( -- * Interface
    checkProg,
    TypeError (..),
    ErrorCase (..),

    -- * Extensionality
    TypeM,
    bad,
    context,
    Checkable (..),
    lookupVar,
    lookupAliases,
    checkOpWith,

    -- * Checkers
    require,
    requireI,
    requirePrimExp,
    checkSubExp,
    checkCerts,
    checkExp,
    checkStms,
    checkStm,
    checkSlice,
    checkType,
    checkExtType,
    matchExtPat,
    matchExtBranchType,
    argType,
    noArgAliases,
    checkArg,
    checkSOACArrayArgs,
    checkLambda,
    checkBody,
    consume,
    binding,
    alternative,
  )
where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Control.Parallel.Strategies
import Data.Bifunctor (first)
import Data.List (find, intercalate, isPrefixOf, sort)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Text qualified as T
import Futhark.Analysis.Alias
import Futhark.Analysis.PrimExp
import Futhark.Construct (instantiateShapes)
import Futhark.IR.Aliases hiding (lookupAliases)
import Futhark.Util
import Futhark.Util.Pretty (align, docText, indent, ppTuple', pretty, (<+>), (</>))

-- | Information about an error during type checking.  The 'Show'
-- instance for this type produces a human-readable description.
data ErrorCase rep
  = TypeError T.Text
  | UnexpectedType (Exp rep) Type [Type]
  | ReturnTypeError Name [ExtType] [ExtType]
  | DupDefinitionError Name
  | DupParamError Name VName
  | DupPatError VName
  | InvalidPatError (Pat (LetDec (Aliases rep))) [ExtType] (Maybe String)
  | UnknownVariableError VName
  | UnknownFunctionError Name
  | ParameterMismatch (Maybe Name) [Type] [Type]
  | SlicingError Shape Int
  | BadAnnotation String Type Type
  | ReturnAliased Name VName
  | UniqueReturnAliased Name
  | NotAnArray VName Type
  | PermutationError [Int] Int (Maybe VName)

instance (Checkable rep) => Show (ErrorCase rep) where
  show :: ErrorCase rep -> String
show (TypeError Text
msg) =
    String
"Type error:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
msg
  show (UnexpectedType Exp rep
e Type
_ []) =
    String
"Type of expression\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack (Doc (ZonkAny 8) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 8) -> Text) -> Doc (ZonkAny 8) -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc (ZonkAny 8) -> Doc (ZonkAny 8))
-> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall a b. (a -> b) -> a -> b
$ Exp rep -> Doc (ZonkAny 8)
forall a ann. Pretty a => a -> Doc ann
forall ann. Exp rep -> Doc ann
pretty Exp rep
e)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\ncannot have any type - possibly a bug in the type checker."
  show (UnexpectedType Exp rep
e Type
t [Type]
ts) =
    String
"Type of expression\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack (Doc (ZonkAny 9) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 9) -> Text) -> Doc (ZonkAny 9) -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Doc (ZonkAny 9) -> Doc (ZonkAny 9)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc (ZonkAny 9) -> Doc (ZonkAny 9))
-> Doc (ZonkAny 9) -> Doc (ZonkAny 9)
forall a b. (a -> b) -> a -> b
$ Exp rep -> Doc (ZonkAny 9)
forall a ann. Pretty a => a -> Doc ann
forall ann. Exp rep -> Doc ann
pretty Exp rep
e)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nmust be one of "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ((Type -> String) -> [Type] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Type -> String
forall a. Pretty a => a -> String
prettyString [Type]
ts)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", but is "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Pretty a => a -> String
prettyString Type
t
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
  show (ReturnTypeError Name
fname [ExtType]
rettype [ExtType]
bodytype) =
    String
"Declaration of function "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
fname
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" declares return type\n  "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack ([ExtType] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [ExtType]
rettype)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nBut body has type\n  "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack ([ExtType] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [ExtType]
bodytype)
  show (DupDefinitionError Name
name) =
    String
"Duplicate definition of function " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
name
  show (DupParamError Name
funname VName
paramname) =
    String
"Parameter "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
paramname
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" mentioned multiple times in argument list of function "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
funname
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
  show (DupPatError VName
name) =
    String
"Variable " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" bound twice in pattern."
  show (InvalidPatError Pat (LetDec (Aliases rep))
pat [ExtType]
t Maybe String
desc) =
    String
"Pat\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Pat (VarAliases, LetDec rep) -> String
forall a. Pretty a => a -> String
prettyString Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\ncannot match value of type\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack ([ExtType] -> Text
forall a. Pretty a => [a] -> Text
prettyTupleLines [ExtType]
t)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
end
    where
      end :: String
end = case Maybe String
desc of
        Maybe String
Nothing -> String
"."
        Just String
desc' -> String
":\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
desc'
  show (UnknownVariableError VName
name) =
    String
"Use of unknown variable " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
  show (UnknownFunctionError Name
fname) =
    String
"Call of unknown function " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
fname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
  show (ParameterMismatch Maybe Name
fname [Type]
expected [Type]
got) =
    String
"In call of "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
fname'
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
":\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"expecting "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
nexpected
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" arguments of type(s)\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ((Type -> String) -> [Type] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Type -> String
forall a. Pretty a => a -> String
prettyString [Type]
expected)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nGot "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
ngot
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" arguments of types\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ((Type -> String) -> [Type] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Type -> String
forall a. Pretty a => a -> String
prettyString [Type]
got)
    where
      nexpected :: Int
nexpected = [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
expected
      ngot :: Int
ngot = [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
got
      fname' :: String
fname' = String -> (Name -> String) -> Maybe Name -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe String
"anonymous function" ((String
"function " String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameToString) Maybe Name
fname
  show (SlicingError Shape
dims Int
got) =
    Int -> String
forall a. Show a => a -> String
show Int
got String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" indices given, but type of indexee has shape " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Shape -> String
forall a. Pretty a => a -> String
prettyString Shape
dims
  show (BadAnnotation String
desc Type
expected Type
got) =
    String
"Annotation of \""
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
desc
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\" type of expression is "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Pretty a => a -> String
prettyString Type
expected
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", but derived to be "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Pretty a => a -> String
prettyString Type
got
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
  show (ReturnAliased Name
fname VName
name) =
    String
"Unique return value of function "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
fname
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" is aliased to "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
name
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", which is not consumed."
  show (UniqueReturnAliased Name
fname) =
    String
"A unique tuple element of return value of function "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
fname
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" is aliased to some other tuple component."
  show (NotAnArray VName
e Type
t) =
    String
"The expression "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
e
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" is expected to be an array, but is "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Pretty a => a -> String
prettyString Type
t
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
  show (PermutationError [Int]
perm Int
rank Maybe VName
name) =
    String
"The permutation ("
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ((Int -> String) -> [Int] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Int -> String
forall a. Show a => a -> String
show [Int]
perm)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") is not valid for array "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
name'
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"of rank "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
rank
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
    where
      name' :: String
name' = String -> (VName -> String) -> Maybe VName -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe String
"" ((String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" ") ShowS -> (VName -> String) -> VName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> String
forall a. Pretty a => a -> String
prettyString) Maybe VName
name

-- | A type error.
data TypeError rep = Error [T.Text] (ErrorCase rep)

instance (Checkable rep) => Show (TypeError rep) where
  show :: TypeError rep -> String
show (Error [] ErrorCase rep
err) =
    ErrorCase rep -> String
forall a. Show a => a -> String
show ErrorCase rep
err
  show (Error [Text]
msgs ErrorCase rep
err) =
    String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"\n" ((Text -> String) -> [Text] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Text -> String
T.unpack [Text]
msgs) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ ErrorCase rep -> String
forall a. Show a => a -> String
show ErrorCase rep
err

-- | A tuple of a return type and a list of parameters, possibly
-- named.
type FunBinding rep = ([(RetType (Aliases rep), RetAls)], [FParam (Aliases rep)])

type VarBinding rep = NameInfo (Aliases rep)

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

data Occurence = Occurence
  { Occurence -> Names
observed :: Names,
    Occurence -> Names
consumed :: Names
  }
  deriving (Occurence -> Occurence -> Bool
(Occurence -> Occurence -> Bool)
-> (Occurence -> Occurence -> Bool) -> Eq Occurence
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Occurence -> Occurence -> Bool
== :: Occurence -> Occurence -> Bool
$c/= :: Occurence -> Occurence -> Bool
/= :: Occurence -> Occurence -> Bool
Eq, Int -> Occurence -> ShowS
[Occurence] -> ShowS
Occurence -> String
(Int -> Occurence -> ShowS)
-> (Occurence -> String)
-> ([Occurence] -> ShowS)
-> Show Occurence
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Occurence -> ShowS
showsPrec :: Int -> Occurence -> ShowS
$cshow :: Occurence -> String
show :: Occurence -> String
$cshowList :: [Occurence] -> ShowS
showList :: [Occurence] -> ShowS
Show)

observation :: Names -> Occurence
observation :: Names -> Occurence
observation = (Names -> Names -> Occurence) -> Names -> Names -> Occurence
forall a b c. (a -> b -> c) -> b -> a -> c
flip Names -> Names -> Occurence
Occurence Names
forall a. Monoid a => a
mempty

consumption :: Names -> Occurence
consumption :: Names -> Occurence
consumption = Names -> Names -> Occurence
Occurence Names
forall a. Monoid a => a
mempty

nullOccurence :: Occurence -> Bool
nullOccurence :: Occurence -> Bool
nullOccurence Occurence
occ = Occurence -> Names
observed Occurence
occ Names -> Names -> Bool
forall a. Eq a => a -> a -> Bool
== Names
forall a. Monoid a => a
mempty Bool -> Bool -> Bool
&& Occurence -> Names
consumed Occurence
occ Names -> Names -> Bool
forall a. Eq a => a -> a -> Bool
== Names
forall a. Monoid a => a
mempty

type Occurences = [Occurence]

allConsumed :: Occurences -> Names
allConsumed :: [Occurence] -> Names
allConsumed = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names)
-> ([Occurence] -> [Names]) -> [Occurence] -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Occurence -> Names) -> [Occurence] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Names
consumed

seqOccurences :: Occurences -> Occurences -> Occurences
seqOccurences :: [Occurence] -> [Occurence] -> [Occurence]
seqOccurences [Occurence]
occurs1 [Occurence]
occurs2 =
  (Occurence -> Bool) -> [Occurence] -> [Occurence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurence -> Bool) -> Occurence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurence -> Bool
nullOccurence) ((Occurence -> Occurence) -> [Occurence] -> [Occurence]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Occurence
filt [Occurence]
occurs1) [Occurence] -> [Occurence] -> [Occurence]
forall a. [a] -> [a] -> [a]
++ [Occurence]
occurs2
  where
    filt :: Occurence -> Occurence
filt Occurence
occ =
      Occurence
occ {observed = observed occ `namesSubtract` postcons}
    postcons :: Names
postcons = [Occurence] -> Names
allConsumed [Occurence]
occurs2

altOccurences :: Occurences -> Occurences -> Occurences
altOccurences :: [Occurence] -> [Occurence] -> [Occurence]
altOccurences [Occurence]
occurs1 [Occurence]
occurs2 =
  (Occurence -> Bool) -> [Occurence] -> [Occurence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurence -> Bool) -> Occurence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurence -> Bool
nullOccurence) ((Occurence -> Occurence) -> [Occurence] -> [Occurence]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Occurence
filt [Occurence]
occurs1) [Occurence] -> [Occurence] -> [Occurence]
forall a. [a] -> [a] -> [a]
++ [Occurence]
occurs2
  where
    filt :: Occurence -> Occurence
filt Occurence
occ =
      Occurence
occ
        { consumed = consumed occ `namesSubtract` postcons,
          observed = observed occ `namesSubtract` postcons
        }
    postcons :: Names
postcons = [Occurence] -> Names
allConsumed [Occurence]
occurs2

unOccur :: Names -> Occurences -> Occurences
unOccur :: Names -> [Occurence] -> [Occurence]
unOccur Names
to_be_removed = (Occurence -> Bool) -> [Occurence] -> [Occurence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurence -> Bool) -> Occurence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurence -> Bool
nullOccurence) ([Occurence] -> [Occurence])
-> ([Occurence] -> [Occurence]) -> [Occurence] -> [Occurence]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Occurence -> Occurence) -> [Occurence] -> [Occurence]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Occurence
unOccur'
  where
    unOccur' :: Occurence -> Occurence
unOccur' Occurence
occ =
      Occurence
occ
        { observed = observed occ `namesSubtract` to_be_removed,
          consumed = consumed occ `namesSubtract` to_be_removed
        }

-- | The 'Consumption' data structure is used to keep track of which
-- variables have been consumed, as well as whether a violation has been detected.
data Consumption
  = ConsumptionError T.Text
  | Consumption Occurences
  deriving (Int -> Consumption -> ShowS
[Consumption] -> ShowS
Consumption -> String
(Int -> Consumption -> ShowS)
-> (Consumption -> String)
-> ([Consumption] -> ShowS)
-> Show Consumption
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Consumption -> ShowS
showsPrec :: Int -> Consumption -> ShowS
$cshow :: Consumption -> String
show :: Consumption -> String
$cshowList :: [Consumption] -> ShowS
showList :: [Consumption] -> ShowS
Show)

instance Semigroup Consumption where
  ConsumptionError Text
e <> :: Consumption -> Consumption -> Consumption
<> Consumption
_ = Text -> Consumption
ConsumptionError Text
e
  Consumption
_ <> ConsumptionError Text
e = Text -> Consumption
ConsumptionError Text
e
  Consumption [Occurence]
o1 <> Consumption [Occurence]
o2
    | VName
v : [VName]
_ <- Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Names
consumed_in_o1 Names -> Names -> Names
`namesIntersection` Names
used_in_o2 =
        Text -> Consumption
ConsumptionError (Text -> Consumption) -> Text -> Consumption
forall a b. (a -> b) -> a -> b
$ Text
"Variable " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" referenced after being consumed."
    | Bool
otherwise =
        [Occurence] -> Consumption
Consumption ([Occurence] -> Consumption) -> [Occurence] -> Consumption
forall a b. (a -> b) -> a -> b
$ [Occurence]
o1 [Occurence] -> [Occurence] -> [Occurence]
`seqOccurences` [Occurence]
o2
    where
      consumed_in_o1 :: Names
consumed_in_o1 = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$ (Occurence -> Names) -> [Occurence] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Names
consumed [Occurence]
o1
      used_in_o2 :: Names
used_in_o2 = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$ (Occurence -> Names) -> [Occurence] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Names
consumed [Occurence]
o2 [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> (Occurence -> Names) -> [Occurence] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Occurence -> Names
observed [Occurence]
o2

instance Monoid Consumption where
  mempty :: Consumption
mempty = [Occurence] -> Consumption
Consumption [Occurence]
forall a. Monoid a => a
mempty

-- | The environment contains a variable table and a function table.
-- Type checking happens with access to this environment.  The
-- function table is only initialised at the very beginning, but the
-- variable table will be extended during type-checking when
-- let-expressions are encountered.
data Env rep = Env
  { forall rep. Env rep -> Map VName (VarBinding rep)
envVtable :: M.Map VName (VarBinding rep),
    forall rep. Env rep -> Map Name (FunBinding rep)
envFtable :: M.Map Name (FunBinding rep),
    forall rep. Env rep -> Op (Aliases rep) -> TypeM rep ()
envCheckOp :: Op (Aliases rep) -> TypeM rep (),
    forall rep. Env rep -> [Text]
envContext :: [T.Text]
  }

data TState = TState
  { TState -> Names
stateNames :: Names,
    TState -> Consumption
stateCons :: Consumption
  }

-- | The type checker runs in this monad.
newtype TypeM rep a
  = TypeM (ReaderT (Env rep) (StateT TState (Either (TypeError rep))) a)
  deriving
    ( Applicative (TypeM rep)
Applicative (TypeM rep) =>
(forall a b. TypeM rep a -> (a -> TypeM rep b) -> TypeM rep b)
-> (forall a b. TypeM rep a -> TypeM rep b -> TypeM rep b)
-> (forall a. a -> TypeM rep a)
-> Monad (TypeM rep)
forall rep. Applicative (TypeM rep)
forall a. a -> TypeM rep a
forall rep a. a -> TypeM rep a
forall a b. TypeM rep a -> TypeM rep b -> TypeM rep b
forall a b. TypeM rep a -> (a -> TypeM rep b) -> TypeM rep b
forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep b
forall rep a b. TypeM rep a -> (a -> TypeM rep b) -> TypeM rep b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall rep a b. TypeM rep a -> (a -> TypeM rep b) -> TypeM rep b
>>= :: forall a b. TypeM rep a -> (a -> TypeM rep b) -> TypeM rep b
$c>> :: forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep b
>> :: forall a b. TypeM rep a -> TypeM rep b -> TypeM rep b
$creturn :: forall rep a. a -> TypeM rep a
return :: forall a. a -> TypeM rep a
Monad,
      (forall a b. (a -> b) -> TypeM rep a -> TypeM rep b)
-> (forall a b. a -> TypeM rep b -> TypeM rep a)
-> Functor (TypeM rep)
forall a b. a -> TypeM rep b -> TypeM rep a
forall a b. (a -> b) -> TypeM rep a -> TypeM rep b
forall rep a b. a -> TypeM rep b -> TypeM rep a
forall rep a b. (a -> b) -> TypeM rep a -> TypeM rep b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall rep a b. (a -> b) -> TypeM rep a -> TypeM rep b
fmap :: forall a b. (a -> b) -> TypeM rep a -> TypeM rep b
$c<$ :: forall rep a b. a -> TypeM rep b -> TypeM rep a
<$ :: forall a b. a -> TypeM rep b -> TypeM rep a
Functor,
      Functor (TypeM rep)
Functor (TypeM rep) =>
(forall a. a -> TypeM rep a)
-> (forall a b. TypeM rep (a -> b) -> TypeM rep a -> TypeM rep b)
-> (forall a b c.
    (a -> b -> c) -> TypeM rep a -> TypeM rep b -> TypeM rep c)
-> (forall a b. TypeM rep a -> TypeM rep b -> TypeM rep b)
-> (forall a b. TypeM rep a -> TypeM rep b -> TypeM rep a)
-> Applicative (TypeM rep)
forall rep. Functor (TypeM rep)
forall a. a -> TypeM rep a
forall rep a. a -> TypeM rep a
forall a b. TypeM rep a -> TypeM rep b -> TypeM rep a
forall a b. TypeM rep a -> TypeM rep b -> TypeM rep b
forall a b. TypeM rep (a -> b) -> TypeM rep a -> TypeM rep b
forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep a
forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep b
forall rep a b. TypeM rep (a -> b) -> TypeM rep a -> TypeM rep b
forall a b c.
(a -> b -> c) -> TypeM rep a -> TypeM rep b -> TypeM rep c
forall rep a b c.
(a -> b -> c) -> TypeM rep a -> TypeM rep b -> TypeM rep c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall rep a. a -> TypeM rep a
pure :: forall a. a -> TypeM rep a
$c<*> :: forall rep a b. TypeM rep (a -> b) -> TypeM rep a -> TypeM rep b
<*> :: forall a b. TypeM rep (a -> b) -> TypeM rep a -> TypeM rep b
$cliftA2 :: forall rep a b c.
(a -> b -> c) -> TypeM rep a -> TypeM rep b -> TypeM rep c
liftA2 :: forall a b c.
(a -> b -> c) -> TypeM rep a -> TypeM rep b -> TypeM rep c
$c*> :: forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep b
*> :: forall a b. TypeM rep a -> TypeM rep b -> TypeM rep b
$c<* :: forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep a
<* :: forall a b. TypeM rep a -> TypeM rep b -> TypeM rep a
Applicative,
      MonadReader (Env rep),
      MonadState TState
    )

instance
  (Checkable rep) =>
  HasScope (Aliases rep) (TypeM rep)
  where
  lookupType :: VName -> TypeM rep Type
lookupType = (NameInfo (Aliases rep) -> Type)
-> TypeM rep (NameInfo (Aliases rep)) -> TypeM rep Type
forall a b. (a -> b) -> TypeM rep a -> TypeM rep b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NameInfo (Aliases rep) -> Type
forall t. Typed t => t -> Type
typeOf (TypeM rep (NameInfo (Aliases rep)) -> TypeM rep Type)
-> (VName -> TypeM rep (NameInfo (Aliases rep)))
-> VName
-> TypeM rep Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> TypeM rep (NameInfo (Aliases rep))
forall rep. VName -> TypeM rep (NameInfo (Aliases rep))
lookupVar
  askScope :: TypeM rep (Scope (Aliases rep))
askScope = (Env rep -> Scope (Aliases rep)) -> TypeM rep (Scope (Aliases rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env rep -> Scope (Aliases rep))
 -> TypeM rep (Scope (Aliases rep)))
-> (Env rep -> Scope (Aliases rep))
-> TypeM rep (Scope (Aliases rep))
forall a b. (a -> b) -> a -> b
$ [(VName, NameInfo (Aliases rep))] -> Scope (Aliases rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, NameInfo (Aliases rep))] -> Scope (Aliases rep))
-> (Env rep -> [(VName, NameInfo (Aliases rep))])
-> Env rep
-> Scope (Aliases rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, NameInfo (Aliases rep))
 -> Maybe (VName, NameInfo (Aliases rep)))
-> [(VName, NameInfo (Aliases rep))]
-> [(VName, NameInfo (Aliases rep))]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName, NameInfo (Aliases rep))
-> Maybe (VName, NameInfo (Aliases rep))
forall {a} {b}. (a, b) -> Maybe (a, b)
varType ([(VName, NameInfo (Aliases rep))]
 -> [(VName, NameInfo (Aliases rep))])
-> (Env rep -> [(VName, NameInfo (Aliases rep))])
-> Env rep
-> [(VName, NameInfo (Aliases rep))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope (Aliases rep) -> [(VName, NameInfo (Aliases rep))]
forall k a. Map k a -> [(k, a)]
M.toList (Scope (Aliases rep) -> [(VName, NameInfo (Aliases rep))])
-> (Env rep -> Scope (Aliases rep))
-> Env rep
-> [(VName, NameInfo (Aliases rep))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env rep -> Scope (Aliases rep)
forall rep. Env rep -> Map VName (VarBinding rep)
envVtable
    where
      varType :: (a, b) -> Maybe (a, b)
varType (a
name, b
dec) = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a
name, b
dec)

runTypeM ::
  Env rep ->
  TypeM rep a ->
  Either (TypeError rep) a
runTypeM :: forall rep a. Env rep -> TypeM rep a -> Either (TypeError rep) a
runTypeM Env rep
env (TypeM ReaderT (Env rep) (StateT TState (Either (TypeError rep))) a
m) =
  StateT TState (Either (TypeError rep)) a
-> TState -> Either (TypeError rep) a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (ReaderT (Env rep) (StateT TState (Either (TypeError rep))) a
-> Env rep -> StateT TState (Either (TypeError rep)) a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Env rep) (StateT TState (Either (TypeError rep))) a
m Env rep
env) (Names -> Consumption -> TState
TState Names
forall a. Monoid a => a
mempty Consumption
forall a. Monoid a => a
mempty)

-- | Signal a type error.
bad :: ErrorCase rep -> TypeM rep a
bad :: forall rep a. ErrorCase rep -> TypeM rep a
bad ErrorCase rep
e = do
  messages <- (Env rep -> [Text]) -> TypeM rep [Text]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env rep -> [Text]
forall rep. Env rep -> [Text]
envContext
  TypeM $ lift $ lift $ Left $ Error (reverse messages) e

tell :: Consumption -> TypeM rep ()
tell :: forall rep. Consumption -> TypeM rep ()
tell Consumption
cons = (TState -> TState) -> TypeM rep ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TState -> TState) -> TypeM rep ())
-> (TState -> TState) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \TState
s -> TState
s {stateCons = stateCons s <> cons}

-- | Add information about what is being type-checked to the current
-- context.  Liberal use of this combinator makes it easier to track
-- type errors, as the strings are added to type errors signalled via
-- 'bad'.
context ::
  T.Text ->
  TypeM rep a ->
  TypeM rep a
context :: forall rep a. Text -> TypeM rep a -> TypeM rep a
context Text
s = (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall a. (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env rep -> Env rep) -> TypeM rep a -> TypeM rep a)
-> (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall a b. (a -> b) -> a -> b
$ \Env rep
env -> Env rep
env {envContext = s : envContext env}

message :: (Pretty a) => T.Text -> a -> T.Text
message :: forall a. Pretty a => Text -> a -> Text
message Text
s a
x = Doc (ZonkAny 0) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 0) -> Text) -> Doc (ZonkAny 0) -> Text
forall a b. (a -> b) -> a -> b
$ Text -> Doc (ZonkAny 0)
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Text
s Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall ann. Doc ann -> Doc ann
align (a -> Doc (ZonkAny 0)
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
x)

-- | Mark a name as bound.  If the name has been bound previously in
-- the program, report a type error.
bound :: VName -> TypeM rep ()
bound :: forall rep. VName -> TypeM rep ()
bound VName
name = do
  already_seen <- (TState -> Bool) -> TypeM rep Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((TState -> Bool) -> TypeM rep Bool)
-> (TState -> Bool) -> TypeM rep Bool
forall a b. (a -> b) -> a -> b
$ VName -> Names -> Bool
nameIn VName
name (Names -> Bool) -> (TState -> Names) -> TState -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TState -> Names
stateNames
  when already_seen . bad . TypeError $
    "Name " <> prettyText name <> " bound twice"
  modify $ \TState
s -> TState
s {stateNames = oneName name <> stateNames s}

occur :: Occurences -> TypeM rep ()
occur :: forall rep. [Occurence] -> TypeM rep ()
occur = Consumption -> TypeM rep ()
forall rep. Consumption -> TypeM rep ()
tell (Consumption -> TypeM rep ())
-> ([Occurence] -> Consumption) -> [Occurence] -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Occurence] -> Consumption
Consumption ([Occurence] -> Consumption)
-> ([Occurence] -> [Occurence]) -> [Occurence] -> Consumption
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Occurence -> Bool) -> [Occurence] -> [Occurence]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Occurence -> Bool) -> Occurence -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Occurence -> Bool
nullOccurence)

-- | Proclaim that we have made read-only use of the given variable.
-- No-op unless the variable is array-typed.
observe ::
  (Checkable rep) =>
  VName ->
  TypeM rep ()
observe :: forall rep. Checkable rep => VName -> TypeM rep ()
observe VName
name = do
  dec <- VName -> TypeM rep (NameInfo (Aliases rep))
forall rep. VName -> TypeM rep (NameInfo (Aliases rep))
lookupVar VName
name
  unless (primType $ typeOf dec) $
    occur [observation $ oneName name <> aliases dec]

-- | Proclaim that we have written to the given variables.
consume :: (Checkable rep) => Names -> TypeM rep ()
consume :: forall rep. Checkable rep => Names -> TypeM rep ()
consume Names
als = do
  scope <- TypeM rep (Scope (Aliases rep))
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let isArray = Bool
-> (NameInfo (Aliases rep) -> Bool)
-> Maybe (NameInfo (Aliases rep))
-> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Bool -> Bool
not (Bool -> Bool)
-> (NameInfo (Aliases rep) -> Bool)
-> NameInfo (Aliases rep)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool)
-> (NameInfo (Aliases rep) -> Type)
-> NameInfo (Aliases rep)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameInfo (Aliases rep) -> Type
forall t. Typed t => t -> Type
typeOf) (Maybe (NameInfo (Aliases rep)) -> Bool)
-> (VName -> Maybe (NameInfo (Aliases rep))) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Scope (Aliases rep) -> Maybe (NameInfo (Aliases rep))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Scope (Aliases rep)
scope)
  occur [consumption $ namesFromList $ filter isArray $ namesToList als]

collectOccurences :: TypeM rep a -> TypeM rep (a, Occurences)
collectOccurences :: forall rep a. TypeM rep a -> TypeM rep (a, [Occurence])
collectOccurences TypeM rep a
m = do
  old <- (TState -> Consumption) -> TypeM rep Consumption
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TState -> Consumption
stateCons
  modify $ \TState
s -> TState
s {stateCons = mempty}
  x <- m
  new <- gets stateCons
  modify $ \TState
s -> TState
s {stateCons = old}
  o <- checkConsumption new
  pure (x, o)

checkOpWith ::
  (Op (Aliases rep) -> TypeM rep ()) ->
  TypeM rep a ->
  TypeM rep a
checkOpWith :: forall rep a.
(Op (Aliases rep) -> TypeM rep ()) -> TypeM rep a -> TypeM rep a
checkOpWith Op (Aliases rep) -> TypeM rep ()
checker = (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall a. (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env rep -> Env rep) -> TypeM rep a -> TypeM rep a)
-> (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall a b. (a -> b) -> a -> b
$ \Env rep
env -> Env rep
env {envCheckOp = checker}

checkConsumption :: Consumption -> TypeM rep Occurences
checkConsumption :: forall rep. Consumption -> TypeM rep [Occurence]
checkConsumption (ConsumptionError Text
e) = ErrorCase rep -> TypeM rep [Occurence]
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep [Occurence])
-> ErrorCase rep -> TypeM rep [Occurence]
forall a b. (a -> b) -> a -> b
$ Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError Text
e
checkConsumption (Consumption [Occurence]
os) = [Occurence] -> TypeM rep [Occurence]
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Occurence]
os

-- | Type check two mutually exclusive control flow branches.  Think
-- @if@.  This interacts with consumption checking, as it is OK for an
-- array to be consumed in both branches.
alternative :: TypeM rep a -> TypeM rep b -> TypeM rep (a, b)
alternative :: forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep (a, b)
alternative TypeM rep a
m1 TypeM rep b
m2 = do
  (x, os1) <- TypeM rep a -> TypeM rep (a, [Occurence])
forall rep a. TypeM rep a -> TypeM rep (a, [Occurence])
collectOccurences TypeM rep a
m1
  (y, os2) <- collectOccurences m2
  tell $ Consumption $ os1 `altOccurences` os2
  pure (x, y)

alternatives :: [TypeM rep ()] -> TypeM rep ()
alternatives :: forall rep. [TypeM rep ()] -> TypeM rep ()
alternatives [] = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
alternatives (TypeM rep ()
x : [TypeM rep ()]
xs) = TypeM rep ((), ()) -> TypeM rep ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TypeM rep ((), ()) -> TypeM rep ())
-> TypeM rep ((), ()) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ TypeM rep ()
x TypeM rep () -> TypeM rep () -> TypeM rep ((), ())
forall rep a b. TypeM rep a -> TypeM rep b -> TypeM rep (a, b)
`alternative` [TypeM rep ()] -> TypeM rep ()
forall rep. [TypeM rep ()] -> TypeM rep ()
alternatives [TypeM rep ()]
xs

-- | Permit consumption of only the specified names.  If one of these
-- names is consumed, the consumption will be rewritten to be a
-- consumption of the corresponding alias set.  Consumption of
-- anything else will result in a type error.
consumeOnlyParams :: [(VName, Names)] -> TypeM rep a -> TypeM rep a
consumeOnlyParams :: forall rep a. [(VName, Names)] -> TypeM rep a -> TypeM rep a
consumeOnlyParams [(VName, Names)]
consumable TypeM rep a
m = do
  (x, os) <- TypeM rep a -> TypeM rep (a, [Occurence])
forall rep a. TypeM rep a -> TypeM rep (a, [Occurence])
collectOccurences TypeM rep a
m
  tell . Consumption =<< mapM inspect os
  pure x
  where
    inspect :: Occurence -> TypeM rep Occurence
inspect Occurence
o = do
      new_consumed <- [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> TypeM rep [Names] -> TypeM rep Names
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> TypeM rep Names) -> [VName] -> TypeM rep [Names]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> TypeM rep Names
wasConsumed (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Occurence -> Names
consumed Occurence
o)
      pure o {consumed = new_consumed}
    wasConsumed :: VName -> TypeM rep Names
wasConsumed VName
v
      | Just Names
als <- VName -> [(VName, Names)] -> Maybe Names
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VName
v [(VName, Names)]
consumable = Names -> TypeM rep Names
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Names
als
      | Bool
otherwise =
          ErrorCase rep -> TypeM rep Names
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep Names)
-> ([Text] -> ErrorCase rep) -> [Text] -> TypeM rep Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> ErrorCase rep)
-> ([Text] -> Text) -> [Text] -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Text
T.unlines ([Text] -> TypeM rep Names) -> [Text] -> TypeM rep Names
forall a b. (a -> b) -> a -> b
$
            [ VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" was invalidly consumed.",
              Text
what Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" can be consumed here."
            ]
    what :: Text
what
      | [(VName, Names)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(VName, Names)]
consumable = Text
"Nothing"
      | Bool
otherwise = Text
"Only " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " (((VName, Names) -> Text) -> [(VName, Names)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Text
forall a. Pretty a => a -> Text
prettyText (VName -> Text)
-> ((VName, Names) -> VName) -> (VName, Names) -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Names) -> VName
forall a b. (a, b) -> a
fst) [(VName, Names)]
consumable)

-- | Given the immediate aliases, compute the full transitive alias
-- set (including the immediate aliases).
expandAliases :: Names -> Env rep -> Names
expandAliases :: forall rep. Names -> Env rep -> Names
expandAliases Names
names Env rep
env = Names
names Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
aliasesOfAliases
  where
    aliasesOfAliases :: Names
aliasesOfAliases = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> (Names -> [Names]) -> Names -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Names
look ([VName] -> [Names]) -> (Names -> [VName]) -> Names -> [Names]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$ Names
names
    look :: VName -> Names
look VName
k = case VName -> Map VName (VarBinding rep) -> Maybe (VarBinding rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
k (Map VName (VarBinding rep) -> Maybe (VarBinding rep))
-> Map VName (VarBinding rep) -> Maybe (VarBinding rep)
forall a b. (a -> b) -> a -> b
$ Env rep -> Map VName (VarBinding rep)
forall rep. Env rep -> Map VName (VarBinding rep)
envVtable Env rep
env of
      Just (LetName (VarAliases
als, LetDec rep
_)) -> VarAliases -> Names
unAliases VarAliases
als
      Maybe (VarBinding rep)
_ -> Names
forall a. Monoid a => a
mempty

binding ::
  (Checkable rep) =>
  Scope (Aliases rep) ->
  TypeM rep a ->
  TypeM rep a
binding :: forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
binding Scope (Aliases rep)
stms = TypeM rep a -> TypeM rep a
check (TypeM rep a -> TypeM rep a)
-> (TypeM rep a -> TypeM rep a) -> TypeM rep a -> TypeM rep a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall a. (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Env rep -> Scope (Aliases rep) -> Env rep
forall {rep}.
Typed (LetDec rep) =>
Env rep -> Map VName (NameInfo (Aliases rep)) -> Env rep
`bindVars` Scope (Aliases rep)
stms)
  where
    bindVars :: Env rep -> Map VName (NameInfo (Aliases rep)) -> Env rep
bindVars Env rep
orig_env = (Env rep -> VName -> NameInfo (Aliases rep) -> Env rep)
-> Env rep -> Map VName (NameInfo (Aliases rep)) -> Env rep
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' (Env rep -> Env rep -> VName -> NameInfo (Aliases rep) -> Env rep
forall {rep} {rep}.
Typed (LetDec rep) =>
Env rep -> Env rep -> VName -> NameInfo (Aliases rep) -> Env rep
bindVar Env rep
orig_env) Env rep
orig_env
    boundnames :: [VName]
boundnames = Scope (Aliases rep) -> [VName]
forall k a. Map k a -> [k]
M.keys Scope (Aliases rep)
stms

    bindVar :: Env rep -> Env rep -> VName -> NameInfo (Aliases rep) -> Env rep
bindVar Env rep
orig_env Env rep
env VName
name (LetName (AliasDec Names
als, LetDec rep
dec)) =
      let als' :: Names
als'
            | Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (LetDec rep -> Type
forall t. Typed t => t -> Type
typeOf LetDec rep
dec) = Names
forall a. Monoid a => a
mempty
            | Bool
otherwise = Names -> Env rep -> Names
forall rep. Names -> Env rep -> Names
expandAliases Names
als Env rep
orig_env
       in Env rep
env
            { envVtable =
                M.insert name (LetName (AliasDec als', dec)) $ envVtable env
            }
    bindVar Env rep
_ Env rep
env VName
name NameInfo (Aliases rep)
dec =
      Env rep
env {envVtable = M.insert name dec $ envVtable env}

    -- Check whether the bound variables have been used correctly
    -- within their scope.
    check :: TypeM rep a -> TypeM rep a
check TypeM rep a
m = do
      (VName -> TypeM rep ()) -> [VName] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ VName -> TypeM rep ()
forall rep. VName -> TypeM rep ()
bound ([VName] -> TypeM rep ()) -> [VName] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ Scope (Aliases rep) -> [VName]
forall k a. Map k a -> [k]
M.keys Scope (Aliases rep)
stms
      (a, os) <- TypeM rep a -> TypeM rep (a, [Occurence])
forall rep a. TypeM rep a -> TypeM rep (a, [Occurence])
collectOccurences TypeM rep a
m
      tell $ Consumption $ unOccur (namesFromList boundnames) os
      pure a

lookupVar :: VName -> TypeM rep (NameInfo (Aliases rep))
lookupVar :: forall rep. VName -> TypeM rep (NameInfo (Aliases rep))
lookupVar VName
name = do
  stm <- (Env rep -> Maybe (VarBinding rep))
-> TypeM rep (Maybe (VarBinding rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env rep -> Maybe (VarBinding rep))
 -> TypeM rep (Maybe (VarBinding rep)))
-> (Env rep -> Maybe (VarBinding rep))
-> TypeM rep (Maybe (VarBinding rep))
forall a b. (a -> b) -> a -> b
$ VName -> Map VName (VarBinding rep) -> Maybe (VarBinding rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName (VarBinding rep) -> Maybe (VarBinding rep))
-> (Env rep -> Map VName (VarBinding rep))
-> Env rep
-> Maybe (VarBinding rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env rep -> Map VName (VarBinding rep)
forall rep. Env rep -> Map VName (VarBinding rep)
envVtable
  case stm of
    Maybe (VarBinding rep)
Nothing -> ErrorCase rep -> TypeM rep (VarBinding rep)
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep (VarBinding rep))
-> ErrorCase rep -> TypeM rep (VarBinding rep)
forall a b. (a -> b) -> a -> b
$ VName -> ErrorCase rep
forall rep. VName -> ErrorCase rep
UnknownVariableError VName
name
    Just VarBinding rep
dec -> VarBinding rep -> TypeM rep (VarBinding rep)
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VarBinding rep
dec

lookupAliases :: (Checkable rep) => VName -> TypeM rep Names
lookupAliases :: forall rep. Checkable rep => VName -> TypeM rep Names
lookupAliases VName
name = do
  info <- VName -> TypeM rep (NameInfo (Aliases rep))
forall rep. VName -> TypeM rep (NameInfo (Aliases rep))
lookupVar VName
name
  pure $
    if primType $ typeOf info
      then mempty
      else oneName name <> aliases info

aliases :: NameInfo (Aliases rep) -> Names
aliases :: forall rep. NameInfo (Aliases rep) -> Names
aliases (LetName (VarAliases
als, LetDec rep
_)) = VarAliases -> Names
unAliases VarAliases
als
aliases NameInfo (Aliases rep)
_ = Names
forall a. Monoid a => a
mempty

subExpAliasesM :: (Checkable rep) => SubExp -> TypeM rep Names
subExpAliasesM :: forall rep. Checkable rep => SubExp -> TypeM rep Names
subExpAliasesM Constant {} = Names -> TypeM rep Names
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Names
forall a. Monoid a => a
mempty
subExpAliasesM (Var VName
v) = VName -> TypeM rep Names
forall rep. Checkable rep => VName -> TypeM rep Names
lookupAliases VName
v

lookupFun ::
  (Checkable rep) =>
  Name ->
  [SubExp] ->
  TypeM rep ([(RetType rep, RetAls)], [DeclType])
lookupFun :: forall rep.
Checkable rep =>
Name -> [SubExp] -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
lookupFun Name
fname [SubExp]
args = do
  stm <- (Env rep
 -> Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
-> TypeM
     rep (Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env rep
  -> Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
 -> TypeM
      rep (Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)])))
-> (Env rep
    -> Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
-> TypeM
     rep (Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall a b. (a -> b) -> a -> b
$ Name
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
 -> Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
-> (Env rep
    -> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
-> Env rep
-> Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env rep
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
Env rep -> Map Name (FunBinding rep)
forall rep. Env rep -> Map Name (FunBinding rep)
envFtable
  case stm of
    Maybe ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
Nothing -> ErrorCase rep -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ([(RetType rep, RetAls)], [DeclType]))
-> ErrorCase rep -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall a b. (a -> b) -> a -> b
$ Name -> ErrorCase rep
forall rep. Name -> ErrorCase rep
UnknownFunctionError Name
fname
    Just ([(RetType rep, RetAls)]
ftype, [Param (FParamInfo rep)]
params) -> do
      argts <- (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
args
      case applyRetType (map fst ftype) params $ zip args argts of
        Maybe [RetType rep]
Nothing ->
          ErrorCase rep -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ([(RetType rep, RetAls)], [DeclType]))
-> ErrorCase rep -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall a b. (a -> b) -> a -> b
$ Maybe Name -> [Type] -> [Type] -> ErrorCase rep
forall rep. Maybe Name -> [Type] -> [Type] -> ErrorCase rep
ParameterMismatch (Name -> Maybe Name
forall a. a -> Maybe a
Just Name
fname) ((Param (FParamInfo rep) -> Type)
-> [Param (FParamInfo rep)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> Type
forall dec. Typed dec => Param dec -> Type
paramType [Param (FParamInfo rep)]
params) [Type]
argts
        Just [RetType rep]
rt ->
          ([(RetType rep, RetAls)], [DeclType])
-> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([RetType rep] -> [RetAls] -> [(RetType rep, RetAls)]
forall a b. [a] -> [b] -> [(a, b)]
zip [RetType rep]
rt ([RetAls] -> [(RetType rep, RetAls)])
-> [RetAls] -> [(RetType rep, RetAls)]
forall a b. (a -> b) -> a -> b
$ ((RetType rep, RetAls) -> RetAls)
-> [(RetType rep, RetAls)] -> [RetAls]
forall a b. (a -> b) -> [a] -> [b]
map (RetType rep, RetAls) -> RetAls
forall a b. (a, b) -> b
snd [(RetType rep, RetAls)]
ftype, (Param (FParamInfo rep) -> DeclType)
-> [Param (FParamInfo rep)] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType [Param (FParamInfo rep)]
params)

-- | @checkAnnotation loc s t1 t2@ checks if @t2@ is equal to
-- @t1@.  If not, a 'BadAnnotation' is raised.
checkAnnotation ::
  String ->
  Type ->
  Type ->
  TypeM rep ()
checkAnnotation :: forall rep. String -> Type -> Type -> TypeM rep ()
checkAnnotation String
desc Type
t1 Type
t2
  | Type
t2 Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
t1 = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise = ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ String -> Type -> Type -> ErrorCase rep
forall rep. String -> Type -> Type -> ErrorCase rep
BadAnnotation String
desc Type
t1 Type
t2

-- | @require ts se@ causes a '(TypeError vn)' if the type of @se@ is
-- not a subtype of one of the types in @ts@.
require :: (Checkable rep) => [Type] -> SubExp -> TypeM rep ()
require :: forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [Type]
ts SubExp
se = do
  t <- SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
se
  unless (t `elem` ts) $ bad $ UnexpectedType (BasicOp $ SubExp se) t ts

-- | Variant of 'require' working on variable names.
requireI :: (Checkable rep) => [Type] -> VName -> TypeM rep ()
requireI :: forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
requireI [Type]
ts VName
ident = [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [Type]
ts (SubExp -> TypeM rep ()) -> SubExp -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ident

checkArrIdent ::
  (Checkable rep) =>
  VName ->
  TypeM rep (Shape, PrimType)
checkArrIdent :: forall rep. Checkable rep => VName -> TypeM rep (Shape, PrimType)
checkArrIdent VName
v = do
  t <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  case t of
    Array PrimType
pt Shape
shape NoUniqueness
_ -> (Shape, PrimType) -> TypeM rep (Shape, PrimType)
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Shape
shape, PrimType
pt)
    Type
_ -> ErrorCase rep -> TypeM rep (Shape, PrimType)
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep (Shape, PrimType))
-> ErrorCase rep -> TypeM rep (Shape, PrimType)
forall a b. (a -> b) -> a -> b
$ VName -> Type -> ErrorCase rep
forall rep. VName -> Type -> ErrorCase rep
NotAnArray VName
v Type
t

checkAccIdent ::
  (Checkable rep) =>
  VName ->
  TypeM rep (Shape, [Type])
checkAccIdent :: forall rep. Checkable rep => VName -> TypeM rep (Shape, [Type])
checkAccIdent VName
v = do
  t <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  case t of
    Acc VName
_ Shape
ispace [Type]
ts NoUniqueness
_ ->
      (Shape, [Type]) -> TypeM rep (Shape, [Type])
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Shape
ispace, [Type]
ts)
    Type
_ ->
      ErrorCase rep -> TypeM rep (Shape, [Type])
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep (Shape, [Type]))
-> (Text -> ErrorCase rep) -> Text -> TypeM rep (Shape, [Type])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep (Shape, [Type]))
-> Text -> TypeM rep (Shape, [Type])
forall a b. (a -> b) -> a -> b
$
        VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" should be an accumulator but is of type "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Type -> Text
forall a. Pretty a => a -> Text
prettyText Type
t

checkOpaques :: OpaqueTypes -> Either (TypeError rep) ()
checkOpaques :: forall rep. OpaqueTypes -> Either (TypeError rep) ()
checkOpaques (OpaqueTypes [(Name, OpaqueType)]
types) = [Name] -> [(Name, OpaqueType)] -> Either (TypeError rep) ()
forall {rep}.
[Name] -> [(Name, OpaqueType)] -> Either (TypeError rep) ()
descend [] [(Name, OpaqueType)]
types
  where
    descend :: [Name] -> [(Name, OpaqueType)] -> Either (TypeError rep) ()
descend [Name]
_ [] = () -> Either (TypeError rep) ()
forall a. a -> Either (TypeError rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    descend [Name]
known ((Name
name, OpaqueType
t) : [(Name, OpaqueType)]
ts) = do
      [Name] -> OpaqueType -> Either (TypeError rep) ()
forall {t :: * -> *} {rep}.
Foldable t =>
t Name -> OpaqueType -> Either (TypeError rep) ()
check [Name]
known OpaqueType
t
      [Name] -> [(Name, OpaqueType)] -> Either (TypeError rep) ()
descend (Name
name Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name]
known) [(Name, OpaqueType)]
ts
    check :: t Name -> OpaqueType -> Either (TypeError rep) ()
check t Name
known (OpaqueRecord [(Name, EntryPointType)]
fs) =
      ((Name, EntryPointType) -> Either (TypeError rep) ())
-> [(Name, EntryPointType)] -> Either (TypeError rep) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (t Name -> EntryPointType -> Either (TypeError rep) ()
forall {t :: * -> *} {rep}.
Foldable t =>
t Name -> EntryPointType -> Either (TypeError rep) ()
checkEntryPointType t Name
known (EntryPointType -> Either (TypeError rep) ())
-> ((Name, EntryPointType) -> EntryPointType)
-> (Name, EntryPointType)
-> Either (TypeError rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd) [(Name, EntryPointType)]
fs
    check t Name
known (OpaqueSum [ValueType]
_ [(Name, [(EntryPointType, [Int])])]
cs) =
      ((Name, [(EntryPointType, [Int])]) -> Either (TypeError rep) ())
-> [(Name, [(EntryPointType, [Int])])] -> Either (TypeError rep) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (((EntryPointType, [Int]) -> Either (TypeError rep) ())
-> [(EntryPointType, [Int])] -> Either (TypeError rep) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (t Name -> EntryPointType -> Either (TypeError rep) ()
forall {t :: * -> *} {rep}.
Foldable t =>
t Name -> EntryPointType -> Either (TypeError rep) ()
checkEntryPointType t Name
known (EntryPointType -> Either (TypeError rep) ())
-> ((EntryPointType, [Int]) -> EntryPointType)
-> (EntryPointType, [Int])
-> Either (TypeError rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EntryPointType, [Int]) -> EntryPointType
forall a b. (a, b) -> a
fst) ([(EntryPointType, [Int])] -> Either (TypeError rep) ())
-> ((Name, [(EntryPointType, [Int])]) -> [(EntryPointType, [Int])])
-> (Name, [(EntryPointType, [Int])])
-> Either (TypeError rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [(EntryPointType, [Int])]) -> [(EntryPointType, [Int])]
forall a b. (a, b) -> b
snd) [(Name, [(EntryPointType, [Int])])]
cs
    check t Name
known (OpaqueArray Int
_ Name
v [ValueType]
_) =
      t Name -> EntryPointType -> Either (TypeError rep) ()
forall {t :: * -> *} {rep}.
Foldable t =>
t Name -> EntryPointType -> Either (TypeError rep) ()
checkEntryPointType t Name
known (Name -> EntryPointType
TypeOpaque Name
v)
    check t Name
known (OpaqueRecordArray Int
_ Name
v [(Name, EntryPointType)]
fs) = do
      t Name -> EntryPointType -> Either (TypeError rep) ()
forall {t :: * -> *} {rep}.
Foldable t =>
t Name -> EntryPointType -> Either (TypeError rep) ()
checkEntryPointType t Name
known (Name -> EntryPointType
TypeOpaque Name
v)
      ((Name, EntryPointType) -> Either (TypeError rep) ())
-> [(Name, EntryPointType)] -> Either (TypeError rep) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (t Name -> EntryPointType -> Either (TypeError rep) ()
forall {t :: * -> *} {rep}.
Foldable t =>
t Name -> EntryPointType -> Either (TypeError rep) ()
checkEntryPointType t Name
known (EntryPointType -> Either (TypeError rep) ())
-> ((Name, EntryPointType) -> EntryPointType)
-> (Name, EntryPointType)
-> Either (TypeError rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd) [(Name, EntryPointType)]
fs
    check t Name
_ (OpaqueType [ValueType]
_) =
      () -> Either (TypeError rep) ()
forall a. a -> Either (TypeError rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    checkEntryPointType :: t Name -> EntryPointType -> Either (TypeError rep) ()
checkEntryPointType t Name
known (TypeOpaque Name
s) =
      Bool -> Either (TypeError rep) () -> Either (TypeError rep) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Name
s Name -> t Name -> Bool
forall a. Eq a => a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t Name
known) (Either (TypeError rep) () -> Either (TypeError rep) ())
-> Either (TypeError rep) () -> Either (TypeError rep) ()
forall a b. (a -> b) -> a -> b
$
        TypeError rep -> Either (TypeError rep) ()
forall a b. a -> Either a b
Left (TypeError rep -> Either (TypeError rep) ())
-> (Text -> TypeError rep) -> Text -> Either (TypeError rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> ErrorCase rep -> TypeError rep
forall rep. [Text] -> ErrorCase rep -> TypeError rep
Error [] (ErrorCase rep -> TypeError rep)
-> (Text -> ErrorCase rep) -> Text -> TypeError rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> Either (TypeError rep) ())
-> Text -> Either (TypeError rep) ()
forall a b. (a -> b) -> a -> b
$
          Text
"Opaque not defined before first use: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText Name
s
    checkEntryPointType t Name
_ (TypeTransparent ValueType
_) = () -> Either (TypeError rep) ()
forall a. a -> Either (TypeError rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Type check a program containing arbitrary type information,
-- yielding either a type error or a program with complete type
-- information.
checkProg ::
  (Checkable rep) =>
  Prog (Aliases rep) ->
  Either (TypeError rep) ()
checkProg :: forall rep.
Checkable rep =>
Prog (Aliases rep) -> Either (TypeError rep) ()
checkProg (Prog OpaqueTypes
opaques Stms (Aliases rep)
consts [FunDef (Aliases rep)]
funs) = do
  OpaqueTypes -> Either (TypeError rep) ()
forall rep. OpaqueTypes -> Either (TypeError rep) ()
checkOpaques OpaqueTypes
opaques
  let typeenv :: Env rep
typeenv =
        Env
          { envVtable :: Map VName (VarBinding rep)
envVtable = Map VName (VarBinding rep)
forall k a. Map k a
M.empty,
            envFtable :: Map Name (FunBinding rep)
envFtable = Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
Map Name (FunBinding rep)
forall a. Monoid a => a
mempty,
            envContext :: [Text]
envContext = [],
            envCheckOp :: Op (Aliases rep) -> TypeM rep ()
envCheckOp = Op (Aliases rep) -> TypeM rep ()
forall rep. Checkable rep => Op (Aliases rep) -> TypeM rep ()
checkOp
          }
  let const_names :: [VName]
const_names = (Stm (Aliases rep) -> [VName]) -> Stms (Aliases rep) -> [VName]
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Pat (VarAliases, LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (VarAliases, LetDec rep) -> [VName])
-> (Stm (Aliases rep) -> Pat (VarAliases, LetDec rep))
-> Stm (Aliases rep)
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm (Aliases rep) -> Pat (VarAliases, LetDec rep)
Stm (Aliases rep) -> Pat (LetDec (Aliases rep))
forall rep. Stm rep -> Pat (LetDec rep)
stmPat) Stms (Aliases rep)
consts
      onFunction :: Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> Map VName (VarBinding rep)
-> FunDef (Aliases rep)
-> Either (TypeError rep) ()
onFunction Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
ftable Map VName (VarBinding rep)
vtable FunDef (Aliases rep)
fun = Env rep -> TypeM rep () -> Either (TypeError rep) ()
forall rep a. Env rep -> TypeM rep a -> Either (TypeError rep) a
runTypeM Env rep
typeenv (TypeM rep () -> Either (TypeError rep) ())
-> TypeM rep () -> Either (TypeError rep) ()
forall a b. (a -> b) -> a -> b
$ do
        (TState -> TState) -> TypeM rep ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TState -> TState) -> TypeM rep ())
-> (TState -> TState) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \TState
s -> TState
s {stateNames = namesFromList const_names}
        (Env rep -> Env rep) -> TypeM rep () -> TypeM rep ()
forall a. (Env rep -> Env rep) -> TypeM rep a -> TypeM rep a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\Env rep
env -> Env rep
env {envFtable = ftable, envVtable = vtable}) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
          FunDef (Aliases rep) -> TypeM rep ()
forall rep. Checkable rep => FunDef (Aliases rep) -> TypeM rep ()
checkFun FunDef (Aliases rep)
fun
  ftable <-
    Env rep
-> TypeM
     rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
-> Either
     (TypeError rep)
     (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall rep a. Env rep -> TypeM rep a -> Either (TypeError rep) a
runTypeM Env rep
typeenv TypeM
  rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
buildFtable
  vtable <-
    runTypeM typeenv {envFtable = ftable} $ checkStms consts $ asks envVtable
  sequence_ $ parMap rpar (onFunction ftable vtable) funs
  where
    buildFtable :: TypeM
  rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
buildFtable = do
      table <- TypeM
  rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
TypeM rep (Map Name (FunBinding rep))
forall rep. Checkable rep => TypeM rep (Map Name (FunBinding rep))
initialFtable
      foldM expand table funs
    expand :: Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> FunDef rep
-> TypeM
     rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
expand Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
ftable (FunDef Maybe EntryPoint
_ Attrs
_ Name
name [(RetType rep, RetAls)]
ret [Param (FParamInfo rep)]
params Body rep
_)
      | Name
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member Name
name Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
ftable =
          ErrorCase rep
-> TypeM
     rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep
 -> TypeM
      rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])))
-> ErrorCase rep
-> TypeM
     rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall a b. (a -> b) -> a -> b
$ Name -> ErrorCase rep
forall rep. Name -> ErrorCase rep
DupDefinitionError Name
name
      | Bool
otherwise =
          Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> TypeM
     rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
 -> TypeM
      rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])))
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> TypeM
     rep (Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall a b. (a -> b) -> a -> b
$ Name
-> ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
name ([(RetType rep, RetAls)]
ret, [Param (FParamInfo rep)]
params) Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
ftable

initialFtable ::
  (Checkable rep) =>
  TypeM rep (M.Map Name (FunBinding rep))
initialFtable :: forall rep. Checkable rep => TypeM rep (Map Name (FunBinding rep))
initialFtable = ([(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
 -> Map Name (FunBinding rep))
-> TypeM
     rep [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
-> TypeM rep (Map Name (FunBinding rep))
forall a b. (a -> b) -> TypeM rep a -> TypeM rep b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
-> Map Name ([(RetType rep, RetAls)], [Param (FParamInfo rep)])
[(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
-> Map Name (FunBinding rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (TypeM
   rep [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
 -> TypeM rep (Map Name (FunBinding rep)))
-> TypeM
     rep [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
-> TypeM rep (Map Name (FunBinding rep))
forall a b. (a -> b) -> a -> b
$ ((Name, (PrimType, [PrimType]))
 -> TypeM
      rep (Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)])))
-> [(Name, (PrimType, [PrimType]))]
-> TypeM
     rep [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name, (PrimType, [PrimType]))
-> TypeM
     rep (Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))
forall {t :: * -> *} {rep} {a} {a}.
(Traversable t, Checkable rep, IsRetType a) =>
(a, (PrimType, t PrimType))
-> TypeM rep (a, ([(a, RetAls)], t (Param (FParamInfo rep))))
addBuiltin ([(Name, (PrimType, [PrimType]))]
 -> TypeM
      rep [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))])
-> [(Name, (PrimType, [PrimType]))]
-> TypeM
     rep [(Name, ([(RetType rep, RetAls)], [Param (FParamInfo rep)]))]
forall a b. (a -> b) -> a -> b
$ Map Name (PrimType, [PrimType]) -> [(Name, (PrimType, [PrimType]))]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (PrimType, [PrimType])
builtInFunctions
  where
    addBuiltin :: (a, (PrimType, t PrimType))
-> TypeM rep (a, ([(a, RetAls)], t (Param (FParamInfo rep))))
addBuiltin (a
fname, (PrimType
t, t PrimType
ts)) = do
      ps <- (PrimType -> TypeM rep (Param (FParamInfo rep)))
-> t PrimType -> TypeM rep (t (Param (FParamInfo rep)))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM (VName -> PrimType -> TypeM rep (FParam (Aliases rep))
forall rep.
Checkable rep =>
VName -> PrimType -> TypeM rep (FParam (Aliases rep))
primFParam VName
name) t PrimType
ts
      pure (fname, ([(primRetType t, RetAls mempty mempty)], ps))
    name :: VName
name = Name -> Int -> VName
VName (String -> Name
nameFromString String
"x") Int
0

checkFun ::
  (Checkable rep) =>
  FunDef (Aliases rep) ->
  TypeM rep ()
checkFun :: forall rep. Checkable rep => FunDef (Aliases rep) -> TypeM rep ()
checkFun (FunDef Maybe EntryPoint
_ Attrs
_ Name
fname [(RetType (Aliases rep), RetAls)]
rettype [FParam (Aliases rep)]
params Body (Aliases rep)
body) =
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"In function " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText Name
fname)
    (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ (Name, [(DeclExtType, RetAls)], [(VName, NameInfo (Aliases rep))])
-> Maybe [(VName, Names)] -> TypeM rep [Names] -> TypeM rep ()
forall rep.
Checkable rep =>
(Name, [(DeclExtType, RetAls)], [(VName, NameInfo (Aliases rep))])
-> Maybe [(VName, Names)] -> TypeM rep [Names] -> TypeM rep ()
checkFun'
      ( Name
fname,
        ((RetType rep, RetAls) -> (DeclExtType, RetAls))
-> [(RetType rep, RetAls)] -> [(DeclExtType, RetAls)]
forall a b. (a -> b) -> [a] -> [b]
map ((RetType rep -> DeclExtType)
-> (RetType rep, RetAls) -> (DeclExtType, RetAls)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first RetType rep -> DeclExtType
forall t. DeclExtTyped t => t -> DeclExtType
declExtTypeOf) [(RetType rep, RetAls)]
[(RetType (Aliases rep), RetAls)]
rettype,
        [FParam rep] -> [(VName, NameInfo (Aliases rep))]
forall rep. [FParam rep] -> [(VName, NameInfo (Aliases rep))]
funParamsToNameInfos [FParam rep]
[FParam (Aliases rep)]
params
      )
      ([(VName, Names)] -> Maybe [(VName, Names)]
forall a. a -> Maybe a
Just [(VName, Names)]
consumable)
    (TypeM rep [Names] -> TypeM rep ())
-> TypeM rep [Names] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
      [FParam rep] -> TypeM rep ()
forall rep. Checkable rep => [FParam rep] -> TypeM rep ()
checkFunParams [FParam rep]
[FParam (Aliases rep)]
params
      [RetType rep] -> TypeM rep ()
forall rep. Checkable rep => [RetType rep] -> TypeM rep ()
checkRetType ([RetType rep] -> TypeM rep ()) -> [RetType rep] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ ((RetType rep, RetAls) -> RetType rep)
-> [(RetType rep, RetAls)] -> [RetType rep]
forall a b. (a -> b) -> [a] -> [b]
map (RetType rep, RetAls) -> RetType rep
forall a b. (a, b) -> a
fst [(RetType rep, RetAls)]
[(RetType (Aliases rep), RetAls)]
rettype
      Text -> TypeM rep [Names] -> TypeM rep [Names]
forall rep a. Text -> TypeM rep a -> TypeM rep a
context Text
"When checking function body" (TypeM rep [Names] -> TypeM rep [Names])
-> TypeM rep [Names] -> TypeM rep [Names]
forall a b. (a -> b) -> a -> b
$ [(RetType rep, RetAls)] -> Body (Aliases rep) -> TypeM rep [Names]
forall rep.
Checkable rep =>
[(RetType rep, RetAls)] -> Body (Aliases rep) -> TypeM rep [Names]
checkFunBody [(RetType rep, RetAls)]
[(RetType (Aliases rep), RetAls)]
rettype Body (Aliases rep)
body
  where
    consumable :: [(VName, Names)]
consumable =
      [ (FParam rep -> VName
forall dec. Param dec -> VName
paramName FParam rep
param, Names
forall a. Monoid a => a
mempty)
        | FParam rep
param <- [FParam rep]
[FParam (Aliases rep)]
params,
          DeclType -> Bool
forall shape. TypeBase shape Uniqueness -> Bool
unique (DeclType -> Bool) -> DeclType -> Bool
forall a b. (a -> b) -> a -> b
$ FParam rep -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType FParam rep
param
      ]

funParamsToNameInfos ::
  [FParam rep] ->
  [(VName, NameInfo (Aliases rep))]
funParamsToNameInfos :: forall rep. [FParam rep] -> [(VName, NameInfo (Aliases rep))]
funParamsToNameInfos = (Param (FParamInfo rep) -> (VName, NameInfo (Aliases rep)))
-> [Param (FParamInfo rep)] -> [(VName, NameInfo (Aliases rep))]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> (VName, NameInfo (Aliases rep))
Param (FParamInfo (Aliases rep)) -> (VName, NameInfo (Aliases rep))
forall {rep}. Param (FParamInfo rep) -> (VName, NameInfo rep)
nameTypeAndDec
  where
    nameTypeAndDec :: Param (FParamInfo rep) -> (VName, NameInfo rep)
nameTypeAndDec Param (FParamInfo rep)
fparam =
      ( Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
fparam,
        FParamInfo rep -> NameInfo rep
forall rep. FParamInfo rep -> NameInfo rep
FParamName (FParamInfo rep -> NameInfo rep) -> FParamInfo rep -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ Param (FParamInfo rep) -> FParamInfo rep
forall dec. Param dec -> dec
paramDec Param (FParamInfo rep)
fparam
      )

checkFunParams ::
  (Checkable rep) =>
  [FParam rep] ->
  TypeM rep ()
checkFunParams :: forall rep. Checkable rep => [FParam rep] -> TypeM rep ()
checkFunParams = (Param (FParamInfo rep) -> TypeM rep ())
-> [Param (FParamInfo rep)] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Param (FParamInfo rep) -> TypeM rep ())
 -> [Param (FParamInfo rep)] -> TypeM rep ())
-> (Param (FParamInfo rep) -> TypeM rep ())
-> [Param (FParamInfo rep)]
-> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \Param (FParamInfo rep)
param ->
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"In parameter " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Param (FParamInfo rep) -> Text
forall a. Pretty a => a -> Text
prettyText Param (FParamInfo rep)
param) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    VName -> FParamInfo rep -> TypeM rep ()
forall rep.
Checkable rep =>
VName -> FParamInfo rep -> TypeM rep ()
checkFParamDec (Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
param) (Param (FParamInfo rep) -> FParamInfo rep
forall dec. Param dec -> dec
paramDec Param (FParamInfo rep)
param)

checkLambdaParams ::
  (Checkable rep) =>
  [LParam rep] ->
  TypeM rep ()
checkLambdaParams :: forall rep. Checkable rep => [LParam rep] -> TypeM rep ()
checkLambdaParams = (Param (LParamInfo rep) -> TypeM rep ())
-> [Param (LParamInfo rep)] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Param (LParamInfo rep) -> TypeM rep ())
 -> [Param (LParamInfo rep)] -> TypeM rep ())
-> (Param (LParamInfo rep) -> TypeM rep ())
-> [Param (LParamInfo rep)]
-> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \Param (LParamInfo rep)
param ->
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"In parameter " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Param (LParamInfo rep) -> Text
forall a. Pretty a => a -> Text
prettyText Param (LParamInfo rep)
param) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    VName -> LParamInfo rep -> TypeM rep ()
forall rep.
Checkable rep =>
VName -> LParamInfo rep -> TypeM rep ()
checkLParamDec (Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
param) (Param (LParamInfo rep) -> LParamInfo rep
forall dec. Param dec -> dec
paramDec Param (LParamInfo rep)
param)

checkNoDuplicateParams :: Name -> [VName] -> TypeM rep ()
checkNoDuplicateParams :: forall rep. Name -> [VName] -> TypeM rep ()
checkNoDuplicateParams Name
fname = ([VName] -> VName -> TypeM rep [VName])
-> [VName] -> [VName] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ [VName] -> VName -> TypeM rep [VName]
expand []
  where
    expand :: [VName] -> VName -> TypeM rep [VName]
expand [VName]
seen VName
pname
      | Just VName
_ <- (VName -> Bool) -> [VName] -> Maybe VName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
pname) [VName]
seen =
          ErrorCase rep -> TypeM rep [VName]
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep [VName])
-> ErrorCase rep -> TypeM rep [VName]
forall a b. (a -> b) -> a -> b
$ Name -> VName -> ErrorCase rep
forall rep. Name -> VName -> ErrorCase rep
DupParamError Name
fname VName
pname
      | Bool
otherwise =
          [VName] -> TypeM rep [VName]
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName] -> TypeM rep [VName]) -> [VName] -> TypeM rep [VName]
forall a b. (a -> b) -> a -> b
$ VName
pname VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
seen

checkFun' ::
  (Checkable rep) =>
  ( Name,
    [(DeclExtType, RetAls)],
    [(VName, NameInfo (Aliases rep))]
  ) ->
  Maybe [(VName, Names)] ->
  TypeM rep [Names] ->
  TypeM rep ()
checkFun' :: forall rep.
Checkable rep =>
(Name, [(DeclExtType, RetAls)], [(VName, NameInfo (Aliases rep))])
-> Maybe [(VName, Names)] -> TypeM rep [Names] -> TypeM rep ()
checkFun' (Name
fname, [(DeclExtType, RetAls)]
rettype, [(VName, NameInfo (Aliases rep))]
params) Maybe [(VName, Names)]
consumable TypeM rep [Names]
check = do
  Name -> [VName] -> TypeM rep ()
forall rep. Name -> [VName] -> TypeM rep ()
checkNoDuplicateParams Name
fname [VName]
param_names
  Scope (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
binding ([(VName, NameInfo (Aliases rep))] -> Scope (Aliases rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, NameInfo (Aliases rep))]
params) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    (TypeM rep () -> TypeM rep ())
-> ([(VName, Names)] -> TypeM rep () -> TypeM rep ())
-> Maybe [(VName, Names)]
-> TypeM rep ()
-> TypeM rep ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeM rep () -> TypeM rep ()
forall a. a -> a
id [(VName, Names)] -> TypeM rep () -> TypeM rep ()
forall rep a. [(VName, Names)] -> TypeM rep a -> TypeM rep a
consumeOnlyParams Maybe [(VName, Names)]
consumable (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
      body_aliases <- TypeM rep [Names]
check
      context
        ( "When checking the body aliases: "
            <> prettyText (map namesToList body_aliases)
        )
        $ checkReturnAlias body_aliases
  where
    param_names :: [VName]
param_names = ((VName, NameInfo (Aliases rep)) -> VName)
-> [(VName, NameInfo (Aliases rep))] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, NameInfo (Aliases rep)) -> VName
forall a b. (a, b) -> a
fst [(VName, NameInfo (Aliases rep))]
params

    isParam :: VName -> Bool
isParam = (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
param_names)

    unique_names :: Names
unique_names = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ do
      (v, FParamName t) <- [(VName, NameInfo (Aliases rep))]
params
      guard $ unique $ declTypeOf t
      pure v

    allowedArgAliases :: [Int] -> Names
allowedArgAliases [Int]
pals =
      [VName] -> Names
namesFromList ((Int -> VName) -> [Int] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ([VName]
param_names [VName] -> Int -> VName
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
pals) Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
unique_names

    checkReturnAlias :: [Names] -> TypeM rep ()
checkReturnAlias [Names]
retals = ((Int, (DeclExtType, RetAls)) -> Names -> TypeM rep ())
-> [(Int, (DeclExtType, RetAls))] -> [Names] -> TypeM rep ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Int, (DeclExtType, RetAls)) -> Names -> TypeM rep ()
checkRet ([Int] -> [(DeclExtType, RetAls)] -> [(Int, (DeclExtType, RetAls))]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int
0 :: Int) ..] [(DeclExtType, RetAls)]
rettype) [Names]
retals
      where
        comrades :: [(Int, Names, [Int])]
comrades = [Int] -> [Names] -> [[Int]] -> [(Int, Names, [Int])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int
0 ..] [Names]
retals ([[Int]] -> [(Int, Names, [Int])])
-> [[Int]] -> [(Int, Names, [Int])]
forall a b. (a -> b) -> a -> b
$ ((DeclExtType, RetAls) -> [Int])
-> [(DeclExtType, RetAls)] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (RetAls -> [Int]
otherAls (RetAls -> [Int])
-> ((DeclExtType, RetAls) -> RetAls)
-> (DeclExtType, RetAls)
-> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DeclExtType, RetAls) -> RetAls
forall a b. (a, b) -> b
snd) [(DeclExtType, RetAls)]
rettype

        checkRet :: (Int, (DeclExtType, RetAls)) -> Names -> TypeM rep ()
checkRet (Int
i, (Array {}, RetAls [Int]
pals [Int]
rals)) Names
als
          | [VName]
als'' <- (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter VName -> Bool
isParam ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
als',
            Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
als'' =
              ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> ([Text] -> ErrorCase rep) -> [Text] -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> ErrorCase rep)
-> ([Text] -> Text) -> [Text] -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Text
T.unlines ([Text] -> TypeM rep ()) -> [Text] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
                [ [Text] -> Text
T.unwords [Text
"Result", Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
i, Text
"aliases", [VName] -> Text
forall a. Pretty a => a -> Text
prettyText [VName]
als''],
                  [Text] -> Text
T.unwords [Text
"but is only allowed to alias arguments", Names -> Text
forall a. Pretty a => a -> Text
prettyText Names
allowed_args]
                ]
          | ((Int
j, Names
_, [Int]
_) : [(Int, Names, [Int])]
_) <- ((Int, Names, [Int]) -> Bool)
-> [(Int, Names, [Int])] -> [(Int, Names, [Int])]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Names -> [Int] -> (Int, Names, [Int]) -> Bool
forall {a} {t :: * -> *} {t :: * -> *}.
(Eq a, Foldable t, Foldable t) =>
a -> Names -> t a -> (a, Names, t a) -> Bool
isProblem Int
i Names
als' [Int]
rals) [(Int, Names, [Int])]
comrades =
              ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> ([Text] -> ErrorCase rep) -> [Text] -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> ErrorCase rep)
-> ([Text] -> Text) -> [Text] -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Text
T.unlines ([Text] -> TypeM rep ()) -> [Text] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
                [ [Text] -> Text
T.unwords [Text
"Results", Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
i, Text
"and", Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
j, Text
"alias each other"],
                  [Text] -> Text
T.unwords [Text
"but result", Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
i, Text
"only allowed to alias results", [Int] -> Text
forall a. Pretty a => a -> Text
prettyText [Int]
rals],
                  [Names] -> Text
forall a. Pretty a => a -> Text
prettyText [Names]
retals
                ]
          where
            allowed_args :: Names
allowed_args = [Int] -> Names
allowedArgAliases [Int]
pals
            als' :: Names
als' = Names
als Names -> Names -> Names
`namesSubtract` Names
allowed_args
        checkRet (Int, (DeclExtType, RetAls))
_ Names
_ = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

        isProblem :: a -> Names -> t a -> (a, Names, t a) -> Bool
isProblem a
i Names
als t a
rals (a
j, Names
jals, t a
j_rals) =
          a
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
j Bool -> Bool -> Bool
&& a
j a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` t a
rals Bool -> Bool -> Bool
&& a
i a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` t a
j_rals Bool -> Bool -> Bool
&& Names -> Names -> Bool
namesIntersect Names
als Names
jals

checkSubExp :: (Checkable rep) => SubExp -> TypeM rep Type
checkSubExp :: forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp (Constant PrimValue
val) =
  Type -> TypeM rep Type
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> TypeM rep Type) -> Type -> TypeM rep Type
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
checkSubExp (Var VName
ident) = Text -> TypeM rep Type -> TypeM rep Type
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"In subexp " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
ident) (TypeM rep Type -> TypeM rep Type)
-> TypeM rep Type -> TypeM rep Type
forall a b. (a -> b) -> a -> b
$ do
  VName -> TypeM rep ()
forall rep. Checkable rep => VName -> TypeM rep ()
observe VName
ident
  VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ident

checkCerts :: (Checkable rep) => Certs -> TypeM rep ()
checkCerts :: forall rep. Checkable rep => Certs -> TypeM rep ()
checkCerts (Certs [VName]
cs) = (VName -> TypeM rep ()) -> [VName] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> VName -> TypeM rep ()
forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
requireI [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Unit]) [VName]
cs

checkSubExpRes :: (Checkable rep) => SubExpRes -> TypeM rep Type
checkSubExpRes :: forall rep. Checkable rep => SubExpRes -> TypeM rep Type
checkSubExpRes (SubExpRes Certs
cs SubExp
se) = do
  Certs -> TypeM rep ()
forall rep. Checkable rep => Certs -> TypeM rep ()
checkCerts Certs
cs
  SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
se

checkStms ::
  (Checkable rep) =>
  Stms (Aliases rep) ->
  TypeM rep a ->
  TypeM rep a
checkStms :: forall rep a.
Checkable rep =>
Stms (Aliases rep) -> TypeM rep a -> TypeM rep a
checkStms Stms (Aliases rep)
origstms TypeM rep a
m = [Stm (Aliases rep)] -> TypeM rep a
delve ([Stm (Aliases rep)] -> TypeM rep a)
-> [Stm (Aliases rep)] -> TypeM rep a
forall a b. (a -> b) -> a -> b
$ Stms (Aliases rep) -> [Stm (Aliases rep)]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms (Aliases rep)
origstms
  where
    delve :: [Stm (Aliases rep)] -> TypeM rep a
delve (stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
_ Exp (Aliases rep)
e) : [Stm (Aliases rep)]
stms) = do
      Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Doc (ZonkAny 4) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 4) -> Text) -> Doc (ZonkAny 4) -> Text
forall a b. (a -> b) -> a -> b
$ Doc (ZonkAny 4)
"In expression of statement" Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Pat (VarAliases, LetDec rep) -> Doc (ZonkAny 4)
forall a ann. Pretty a => a -> Doc ann
forall ann. Pat (VarAliases, LetDec rep) -> Doc ann
pretty Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat)) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Exp (Aliases rep) -> TypeM rep ()
forall rep. Checkable rep => Exp (Aliases rep) -> TypeM rep ()
checkExp Exp (Aliases rep)
e
      Stm (Aliases rep) -> TypeM rep a -> TypeM rep a
forall rep a.
Checkable rep =>
Stm (Aliases rep) -> TypeM rep a -> TypeM rep a
checkStm Stm (Aliases rep)
stm (TypeM rep a -> TypeM rep a) -> TypeM rep a -> TypeM rep a
forall a b. (a -> b) -> a -> b
$
        [Stm (Aliases rep)] -> TypeM rep a
delve [Stm (Aliases rep)]
stms
    delve [] =
      TypeM rep a
m

checkResult ::
  (Checkable rep) =>
  Result ->
  TypeM rep ()
checkResult :: forall rep. Checkable rep => Result -> TypeM rep ()
checkResult = (SubExpRes -> TypeM rep Type) -> Result -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SubExpRes -> TypeM rep Type
forall rep. Checkable rep => SubExpRes -> TypeM rep Type
checkSubExpRes

checkFunBody ::
  (Checkable rep) =>
  [(RetType rep, RetAls)] ->
  Body (Aliases rep) ->
  TypeM rep [Names]
checkFunBody :: forall rep.
Checkable rep =>
[(RetType rep, RetAls)] -> Body (Aliases rep) -> TypeM rep [Names]
checkFunBody [(RetType rep, RetAls)]
rt (Body (BodyAliasing
_, BodyDec rep
rep) Stms (Aliases rep)
stms Result
res) = do
  BodyDec rep -> TypeM rep ()
forall rep. Checkable rep => BodyDec rep -> TypeM rep ()
checkBodyDec BodyDec rep
rep
  Stms (Aliases rep) -> TypeM rep [Names] -> TypeM rep [Names]
forall rep a.
Checkable rep =>
Stms (Aliases rep) -> TypeM rep a -> TypeM rep a
checkStms Stms (Aliases rep)
stms (TypeM rep [Names] -> TypeM rep [Names])
-> TypeM rep [Names] -> TypeM rep [Names]
forall a b. (a -> b) -> a -> b
$ do
    Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context Text
"When checking body result" (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ Result -> TypeM rep ()
forall rep. Checkable rep => Result -> TypeM rep ()
checkResult Result
res
    Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context Text
"When matching declared return type to result of body" (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
      [RetType rep] -> Result -> TypeM rep ()
forall rep.
Checkable rep =>
[RetType rep] -> Result -> TypeM rep ()
matchReturnType (((RetType rep, RetAls) -> RetType rep)
-> [(RetType rep, RetAls)] -> [RetType rep]
forall a b. (a -> b) -> [a] -> [b]
map (RetType rep, RetAls) -> RetType rep
forall a b. (a, b) -> a
fst [(RetType rep, RetAls)]
rt) Result
res
    (SubExpRes -> TypeM rep Names) -> Result -> TypeM rep [Names]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SubExp -> TypeM rep Names
forall rep. Checkable rep => SubExp -> TypeM rep Names
subExpAliasesM (SubExp -> TypeM rep Names)
-> (SubExpRes -> SubExp) -> SubExpRes -> TypeM rep Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> SubExp
resSubExp) Result
res

checkLambdaBody ::
  (Checkable rep) =>
  [Type] ->
  Body (Aliases rep) ->
  TypeM rep ()
checkLambdaBody :: forall rep.
Checkable rep =>
[Type] -> Body (Aliases rep) -> TypeM rep ()
checkLambdaBody [Type]
ret (Body (BodyAliasing
_, BodyDec rep
rep) Stms (Aliases rep)
stms Result
res) = do
  BodyDec rep -> TypeM rep ()
forall rep. Checkable rep => BodyDec rep -> TypeM rep ()
checkBodyDec BodyDec rep
rep
  Stms (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Stms (Aliases rep) -> TypeM rep a -> TypeM rep a
checkStms Stms (Aliases rep)
stms (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ [Type] -> Result -> TypeM rep ()
forall rep. Checkable rep => [Type] -> Result -> TypeM rep ()
checkLambdaResult [Type]
ret Result
res

checkLambdaResult ::
  (Checkable rep) =>
  [Type] ->
  Result ->
  TypeM rep ()
checkLambdaResult :: forall rep. Checkable rep => [Type] -> Result -> TypeM rep ()
checkLambdaResult [Type]
ts Result
es
  | [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Result -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
es =
      ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Text
"Lambda has return type "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Type] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [Type]
ts
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" describing "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" values, but body returns "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText (Result -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
es)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" values: "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Result -> Text
forall a. Pretty a => [a] -> Text
prettyTuple Result
es
  | Bool
otherwise = [(Type, SubExpRes)]
-> ((Type, SubExpRes) -> TypeM rep ()) -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Type] -> Result -> [(Type, SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Type]
ts Result
es) (((Type, SubExpRes) -> TypeM rep ()) -> TypeM rep ())
-> ((Type, SubExpRes) -> TypeM rep ()) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \(Type
t, SubExpRes
e) -> do
      et <- SubExpRes -> TypeM rep Type
forall rep. Checkable rep => SubExpRes -> TypeM rep Type
checkSubExpRes SubExpRes
e
      unless (et == t) . bad . TypeError $
        "Subexpression "
          <> prettyText e
          <> " has type "
          <> prettyText et
          <> " but expected "
          <> prettyText t

checkBody ::
  (Checkable rep) =>
  Body (Aliases rep) ->
  TypeM rep [Names]
checkBody :: forall rep.
Checkable rep =>
Body (Aliases rep) -> TypeM rep [Names]
checkBody (Body (BodyAliasing
_, BodyDec rep
rep) Stms (Aliases rep)
stms Result
res) = do
  BodyDec rep -> TypeM rep ()
forall rep. Checkable rep => BodyDec rep -> TypeM rep ()
checkBodyDec BodyDec rep
rep
  Stms (Aliases rep) -> TypeM rep [Names] -> TypeM rep [Names]
forall rep a.
Checkable rep =>
Stms (Aliases rep) -> TypeM rep a -> TypeM rep a
checkStms Stms (Aliases rep)
stms (TypeM rep [Names] -> TypeM rep [Names])
-> TypeM rep [Names] -> TypeM rep [Names]
forall a b. (a -> b) -> a -> b
$ do
    Result -> TypeM rep ()
forall rep. Checkable rep => Result -> TypeM rep ()
checkResult Result
res
    (Names -> Names) -> [Names] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Names -> Names -> Names
`namesSubtract` Names
bound_here) ([Names] -> [Names]) -> TypeM rep [Names] -> TypeM rep [Names]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExpRes -> TypeM rep Names) -> Result -> TypeM rep [Names]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SubExp -> TypeM rep Names
forall rep. Checkable rep => SubExp -> TypeM rep Names
subExpAliasesM (SubExp -> TypeM rep Names)
-> (SubExpRes -> SubExp) -> SubExpRes -> TypeM rep Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> SubExp
resSubExp) Result
res
  where
    bound_here :: Names
bound_here = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Map VName (NameInfo (Aliases rep)) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (NameInfo (Aliases rep)) -> [VName])
-> Map VName (NameInfo (Aliases rep)) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stms (Aliases rep) -> Map VName (NameInfo (Aliases rep))
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms (Aliases rep)
stms

-- | Check a slicing operation of an array of the provided type.
checkSlice :: (Checkable rep) => Type -> Slice SubExp -> TypeM rep ()
checkSlice :: forall rep. Checkable rep => Type -> Slice SubExp -> TypeM rep ()
checkSlice Type
vt (Slice [DimIndex SubExp]
idxes) = do
  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
arrayRank Type
vt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex SubExp]
idxes) (TypeM rep () -> TypeM rep ())
-> (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    Shape -> Int -> ErrorCase rep
forall rep. Shape -> Int -> ErrorCase rep
SlicingError (Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
vt) ([DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex SubExp]
idxes)
  (DimIndex SubExp -> TypeM rep (DimIndex ()))
-> [DimIndex SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SubExp -> TypeM rep ())
-> DimIndex SubExp -> TypeM rep (DimIndex ())
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> DimIndex a -> f (DimIndex b)
traverse ((SubExp -> TypeM rep ())
 -> DimIndex SubExp -> TypeM rep (DimIndex ()))
-> (SubExp -> TypeM rep ())
-> DimIndex SubExp
-> TypeM rep (DimIndex ())
forall a b. (a -> b) -> a -> b
$ [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64]) [DimIndex SubExp]
idxes

checkShape :: (Checkable rep) => Shape -> TypeM rep ()
checkShape :: forall rep. Checkable rep => Shape -> TypeM rep ()
checkShape = (SubExp -> TypeM rep ()) -> Shape -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64])

checkBasicOp :: (Checkable rep) => BasicOp -> TypeM rep ()
checkBasicOp :: forall rep. Checkable rep => BasicOp -> TypeM rep ()
checkBasicOp (SubExp SubExp
es) =
  TypeM rep Type -> TypeM rep ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TypeM rep Type -> TypeM rep ()) -> TypeM rep Type -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
es
checkBasicOp (Opaque OpaqueOp
_ SubExp
es) =
  TypeM rep Type -> TypeM rep ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TypeM rep Type -> TypeM rep ()) -> TypeM rep Type -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
es
checkBasicOp ArrayVal {} =
  -- We assume this is never changed, so no need to check it.
  () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkBasicOp (ArrayLit [] Type
_) =
  () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkBasicOp (ArrayLit (SubExp
e : [SubExp]
es') Type
t) = do
  let check :: Type -> SubExp -> TypeM rep ()
check Type
elemt SubExp
eleme = do
        elemet <- SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
eleme
        unless (elemet == elemt) . bad . TypeError $
          prettyText elemet
            <> " is not of expected type "
            <> prettyText elemt
            <> "."
  et <- SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
e

  -- Compare that type with the one given for the array literal.
  checkAnnotation "array-element" t et

  mapM_ (check et) es'
checkBasicOp (UnOp UnOp
op SubExp
e) = [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ UnOp -> PrimType
unOpType UnOp
op] SubExp
e
checkBasicOp (BinOp BinOp
op SubExp
e1 SubExp
e2) = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (BinOp -> PrimType
binOpType BinOp
op) SubExp
e1 SubExp
e2
checkBasicOp (CmpOp CmpOp
op SubExp
e1 SubExp
e2) = CmpOp -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
CmpOp -> SubExp -> SubExp -> TypeM rep ()
checkCmpOp CmpOp
op SubExp
e1 SubExp
e2
checkBasicOp (ConvOp ConvOp
op SubExp
e) = [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ (PrimType, PrimType) -> PrimType
forall a b. (a, b) -> a
fst ((PrimType, PrimType) -> PrimType)
-> (PrimType, PrimType) -> PrimType
forall a b. (a -> b) -> a -> b
$ ConvOp -> (PrimType, PrimType)
convOpType ConvOp
op] SubExp
e
checkBasicOp (Index VName
ident Slice SubExp
slice) = do
  vt <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ident
  observe ident
  checkSlice vt slice
checkBasicOp (Update Safety
_ VName
src Slice SubExp
slice SubExp
se) = do
  (src_shape, src_pt) <- VName -> TypeM rep (Shape, PrimType)
forall rep. Checkable rep => VName -> TypeM rep (Shape, PrimType)
checkArrIdent VName
src

  se_aliases <- subExpAliasesM se
  when (src `nameIn` se_aliases) $
    bad $
      TypeError "The target of an Update must not alias the value to be written."

  checkSlice (arrayOf (Prim src_pt) src_shape NoUniqueness) slice
  require [arrayOf (Prim src_pt) (sliceShape slice) NoUniqueness] se
  consume =<< lookupAliases src
checkBasicOp (FlatIndex VName
ident FlatSlice SubExp
slice) = do
  vt <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ident
  observe ident
  when (arrayRank vt /= 1) $ bad $ SlicingError (arrayShape vt) 1
  checkFlatSlice slice
checkBasicOp (FlatUpdate VName
src FlatSlice SubExp
slice VName
v) = do
  (src_shape, src_pt) <- VName -> TypeM rep (Shape, PrimType)
forall rep. Checkable rep => VName -> TypeM rep (Shape, PrimType)
checkArrIdent VName
src
  when (shapeRank src_shape /= 1) $ bad $ SlicingError src_shape 1

  v_aliases <- lookupAliases v
  when (src `nameIn` v_aliases) $
    bad $
      TypeError "The target of an Update must not alias the value to be written."

  checkFlatSlice slice
  requireI [arrayOf (Prim src_pt) (Shape (flatSliceDims slice)) NoUniqueness] v
  consume =<< lookupAliases src
checkBasicOp (Iota SubExp
e SubExp
x SubExp
s IntType
et) = do
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64] SubExp
e
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
et] SubExp
x
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
et] SubExp
s
checkBasicOp (Replicate (Shape [SubExp]
dims) SubExp
valexp) = do
  (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64]) [SubExp]
dims
  TypeM rep Type -> TypeM rep ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TypeM rep Type -> TypeM rep ()) -> TypeM rep Type -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
valexp
checkBasicOp (Scratch PrimType
_ [SubExp]
shape) =
  Shape -> TypeM rep ()
forall rep. Checkable rep => Shape -> TypeM rep ()
checkShape (Shape -> TypeM rep ()) -> Shape -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
shape
checkBasicOp (Reshape VName
arrexp NewShape SubExp
newshape) = do
  (arr_shape, _) <- VName -> TypeM rep (Shape, PrimType)
forall rep. Checkable rep => VName -> TypeM rep (Shape, PrimType)
checkArrIdent VName
arrexp
  checkShape $ newShape newshape
  spliced_shape <- foldM checkSplice arr_shape $ dimSplices newshape
  when (spliced_shape /= newShape newshape) . bad . TypeError $
    "Splice produces shape "
      <> prettyText spliced_shape
      <> " but annotation is shape "
      <> prettyText (newShape newshape)
  where
    checkSplice :: Shape -> DimSplice SubExp -> TypeM rep Shape
checkSplice Shape
arr_shape sp :: DimSplice SubExp
sp@(DimSplice Int
i Int
k Shape
shape) = do
      Shape -> TypeM rep ()
forall rep. Checkable rep => Shape -> TypeM rep ()
checkShape Shape
shape
      Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
arr_shape) (TypeM rep () -> TypeM rep ())
-> (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Text
"Splice " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> DimSplice SubExp -> Text
forall a. Pretty a => a -> Text
prettyText DimSplice SubExp
sp Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" cannot be applied to shape " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Shape -> Text
forall a. Pretty a => a -> Text
prettyText Shape
arr_shape
      Shape -> TypeM rep Shape
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Shape -> TypeM rep Shape) -> Shape -> TypeM rep Shape
forall a b. (a -> b) -> a -> b
$ Shape -> DimSplice SubExp -> Shape
forall d. ShapeBase d -> DimSplice d -> ShapeBase d
applySplice Shape
arr_shape DimSplice SubExp
sp
checkBasicOp (Rearrange VName
arr [Int]
perm) = do
  arrt <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
  let rank = Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
arrayRank Type
arrt
  when (length perm /= rank || sort perm /= [0 .. rank - 1]) $
    bad $
      PermutationError perm rank $
        Just arr
checkBasicOp (Concat Int
i (VName
arr1exp :| [VName]
arr2exps) SubExp
ressize) = do
  arr1_dims <- Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (Shape -> [SubExp])
-> ((Shape, PrimType) -> Shape) -> (Shape, PrimType) -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Shape, PrimType) -> Shape
forall a b. (a, b) -> a
fst ((Shape, PrimType) -> [SubExp])
-> TypeM rep (Shape, PrimType) -> TypeM rep [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> TypeM rep (Shape, PrimType)
forall rep. Checkable rep => VName -> TypeM rep (Shape, PrimType)
checkArrIdent VName
arr1exp
  arr2s_dims <- map (shapeDims . fst) <$> mapM checkArrIdent arr2exps
  unless (all ((== dropAt i 1 arr1_dims) . dropAt i 1) arr2s_dims) $
    bad $
      TypeError "Types of arguments to concat do not match."
  require [Prim int64] ressize
checkBasicOp (Manifest VName
arr [Int]
perm) =
  BasicOp -> TypeM rep ()
forall rep. Checkable rep => BasicOp -> TypeM rep ()
checkBasicOp (BasicOp -> TypeM rep ()) -> BasicOp -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ VName -> [Int] -> BasicOp
Rearrange VName
arr [Int]
perm -- Basically same thing!
checkBasicOp (Assert SubExp
e (ErrorMsg [ErrorMsgPart SubExp]
parts)) = do
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Bool] SubExp
e
  (ErrorMsgPart SubExp -> TypeM rep ())
-> [ErrorMsgPart SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ErrorMsgPart SubExp -> TypeM rep ()
forall {rep}. Checkable rep => ErrorMsgPart SubExp -> TypeM rep ()
checkPart [ErrorMsgPart SubExp]
parts
  where
    checkPart :: ErrorMsgPart SubExp -> TypeM rep ()
checkPart ErrorString {} = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    checkPart (ErrorVal PrimType
t SubExp
x) = [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t] SubExp
x
checkBasicOp (UpdateAcc Safety
_ VName
acc [SubExp]
is [SubExp]
ses) = do
  (shape, ts) <- VName -> TypeM rep (Shape, [Type])
forall rep. Checkable rep => VName -> TypeM rep (Shape, [Type])
checkAccIdent VName
acc

  unless (length ses == length ts) . bad . TypeError $
    "Accumulator requires "
      <> prettyText (length ts)
      <> " values, but "
      <> prettyText (length ses)
      <> " provided."

  unless (length is == shapeRank shape) $
    bad . TypeError $
      "Accumulator requires "
        <> prettyText (shapeRank shape)
        <> " indices, but "
        <> prettyText (length is)
        <> " provided."

  zipWithM_ require (map pure ts) ses
  consume =<< lookupAliases acc

matchLoopResultExt ::
  (Checkable rep) =>
  [Param DeclType] ->
  Result ->
  TypeM rep ()
matchLoopResultExt :: forall rep.
Checkable rep =>
[Param DeclType] -> Result -> TypeM rep ()
matchLoopResultExt [Param DeclType]
merge Result
loopres = do
  let rettype_ext :: [ExtType]
rettype_ext =
        [VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
paramName [Param DeclType]
merge) ([ExtType] -> [ExtType]) -> [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
          [Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes ([Type] -> [ExtType]) -> [Type] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
            (Param DeclType -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall t. Typed t => t -> Type
typeOf [Param DeclType]
merge

  bodyt <- (SubExpRes -> TypeM rep Type) -> Result -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExpRes -> TypeM rep Type
forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
subExpResType Result
loopres

  case instantiateShapes (fmap resSubExp . (`maybeNth` loopres)) rettype_ext of
    Maybe [Type]
Nothing ->
      ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Name -> [ExtType] -> [ExtType] -> ErrorCase rep
forall rep. Name -> [ExtType] -> [ExtType] -> ErrorCase rep
ReturnTypeError
          (String -> Name
nameFromString String
"<loop body>")
          [ExtType]
rettype_ext
          ([Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes [Type]
bodyt)
    Just [Type]
rettype' ->
      Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Type]
bodyt [Type] -> [Type] -> Bool
forall u shape.
(Ord u, ArrayShape shape) =>
[TypeBase shape u] -> [TypeBase shape u] -> Bool
`subtypesOf` [Type]
rettype') (TypeM rep () -> TypeM rep ())
-> (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Name -> [ExtType] -> [ExtType] -> ErrorCase rep
forall rep. Name -> [ExtType] -> [ExtType] -> ErrorCase rep
ReturnTypeError
          (String -> Name
nameFromString String
"<loop body>")
          ([Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes [Type]
rettype')
          ([Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes [Type]
bodyt)

allowAllAliases :: Int -> Int -> RetAls
allowAllAliases :: Int -> Int -> RetAls
allowAllAliases Int
n Int
m =
  [Int] -> [Int] -> RetAls
RetAls [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

checkExp ::
  (Checkable rep) =>
  Exp (Aliases rep) ->
  TypeM rep ()
checkExp :: forall rep. Checkable rep => Exp (Aliases rep) -> TypeM rep ()
checkExp (BasicOp BasicOp
op) = BasicOp -> TypeM rep ()
forall rep. Checkable rep => BasicOp -> TypeM rep ()
checkBasicOp BasicOp
op
checkExp (Match [SubExp]
ses [Case (Body (Aliases rep))]
cases Body (Aliases rep)
def_case MatchDec (BranchType (Aliases rep))
info) = do
  ses_ts <- (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp [SubExp]
ses
  alternatives $
    context "in body of last case" (checkCaseBody def_case)
      : map (checkCase ses_ts) cases
  where
    checkVal :: TypeBase shape u -> Maybe PrimValue -> Bool
checkVal TypeBase shape u
t (Just PrimValue
v) = PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
Prim (PrimValue -> PrimType
primValueType PrimValue
v) TypeBase shape u -> TypeBase shape u -> Bool
forall a. Eq a => a -> a -> Bool
== TypeBase shape u
t
    checkVal TypeBase shape u
_ Maybe PrimValue
Nothing = Bool
True
    checkCase :: [Type] -> Case (Body (Aliases rep)) -> TypeM rep ()
checkCase [Type]
ses_ts (Case [Maybe PrimValue]
vs Body (Aliases rep)
body) = do
      let ok :: Bool
ok = [Maybe PrimValue] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe PrimValue]
vs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ses_ts Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Type -> Maybe PrimValue -> Bool)
-> [Type] -> [Maybe PrimValue] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> Maybe PrimValue -> Bool
forall {shape} {u}.
(Eq shape, Eq u) =>
TypeBase shape u -> Maybe PrimValue -> Bool
checkVal [Type]
ses_ts [Maybe PrimValue]
vs)
      Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
ok (TypeM rep () -> TypeM rep ())
-> (Doc (ZonkAny 5) -> TypeM rep ())
-> Doc (ZonkAny 5)
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Doc (ZonkAny 5) -> ErrorCase rep)
-> Doc (ZonkAny 5)
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> ErrorCase rep)
-> (Doc (ZonkAny 5) -> Text) -> Doc (ZonkAny 5) -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 5) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 5) -> TypeM rep ())
-> Doc (ZonkAny 5) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Doc (ZonkAny 5)
"Scrutinee"
          Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 5)] -> Doc (ZonkAny 5)
forall a. [Doc a] -> Doc a
ppTuple' ([Doc (ZonkAny 5)] -> Doc (ZonkAny 5))
-> [Doc (ZonkAny 5)] -> Doc (ZonkAny 5)
forall a b. (a -> b) -> a -> b
$ (SubExp -> Doc (ZonkAny 5)) -> [SubExp] -> [Doc (ZonkAny 5)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc (ZonkAny 5)
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty [SubExp]
ses)
          Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 5)
"cannot match pattern"
          Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 5)] -> Doc (ZonkAny 5)
forall a. [Doc a] -> Doc a
ppTuple' ([Doc (ZonkAny 5)] -> Doc (ZonkAny 5))
-> [Doc (ZonkAny 5)] -> Doc (ZonkAny 5)
forall a b. (a -> b) -> a -> b
$ (Maybe PrimValue -> Doc (ZonkAny 5))
-> [Maybe PrimValue] -> [Doc (ZonkAny 5)]
forall a b. (a -> b) -> [a] -> [b]
map Maybe PrimValue -> Doc (ZonkAny 5)
forall ann. Maybe PrimValue -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [Maybe PrimValue]
vs)
      Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"in body of case " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Maybe PrimValue] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [Maybe PrimValue]
vs) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ Body (Aliases rep) -> TypeM rep ()
checkCaseBody Body (Aliases rep)
body
    checkCaseBody :: Body (Aliases rep) -> TypeM rep ()
checkCaseBody Body (Aliases rep)
body = do
      TypeM rep [Names] -> TypeM rep ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TypeM rep [Names] -> TypeM rep ())
-> TypeM rep [Names] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ Body (Aliases rep) -> TypeM rep [Names]
forall rep.
Checkable rep =>
Body (Aliases rep) -> TypeM rep [Names]
checkBody Body (Aliases rep)
body
      [BranchType rep] -> Body (Aliases rep) -> TypeM rep ()
forall rep.
Checkable rep =>
[BranchType rep] -> Body (Aliases rep) -> TypeM rep ()
matchBranchType (MatchDec (BranchType rep) -> [BranchType rep]
forall rt. MatchDec rt -> [rt]
matchReturns MatchDec (BranchType rep)
MatchDec (BranchType (Aliases rep))
info) Body (Aliases rep)
body
checkExp (Apply Name
fname [(SubExp, Diet)]
args [(RetType (Aliases rep), RetAls)]
rettype_annot Safety
_) = do
  (rettype_derived, paramtypes) <- Name -> [SubExp] -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall rep.
Checkable rep =>
Name -> [SubExp] -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
lookupFun Name
fname ([SubExp] -> TypeM rep ([(RetType rep, RetAls)], [DeclType]))
-> [SubExp] -> TypeM rep ([(RetType rep, RetAls)], [DeclType])
forall a b. (a -> b) -> a -> b
$ ((SubExp, Diet) -> SubExp) -> [(SubExp, Diet)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst [(SubExp, Diet)]
args
  argflows <- mapM (checkArg . fst) args
  when (rettype_derived /= rettype_annot) $
    bad . TypeError . docText $
      "Expected apply result type:"
        </> indent 2 (pretty $ map fst rettype_derived)
        </> "But annotation is:"
        </> indent 2 (pretty $ map fst rettype_annot)
  consumeArgs paramtypes argflows
checkExp (Loop [(FParam (Aliases rep), SubExp)]
merge LoopForm
form Body (Aliases rep)
loopbody) = do
  let ([Param (FParamInfo rep)]
mergepat, [SubExp]
mergeexps) = [(Param (FParamInfo rep), SubExp)]
-> ([Param (FParamInfo rep)], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge
  mergeargs <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Arg
forall rep. Checkable rep => SubExp -> TypeM rep Arg
checkArg [SubExp]
mergeexps

  checkLoopArgs

  binding (scopeOfLoopForm form) $ do
    form_consumable <- checkForm mergeargs form

    let rettype = (Param (FParamInfo rep) -> DeclType)
-> [Param (FParamInfo rep)] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType [Param (FParamInfo rep)]
mergepat
        consumable =
          [ (Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
param, Names
forall a. Monoid a => a
mempty)
            | Param (FParamInfo rep)
param <- [Param (FParamInfo rep)]
mergepat,
              DeclType -> Bool
forall shape. TypeBase shape Uniqueness -> Bool
unique (DeclType -> Bool) -> DeclType -> Bool
forall a b. (a -> b) -> a -> b
$ Param (FParamInfo rep) -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType Param (FParamInfo rep)
param
          ]
            [(VName, Names)] -> [(VName, Names)] -> [(VName, Names)]
forall a. [a] -> [a] -> [a]
++ [(VName, Names)]
form_consumable

    context "Inside the loop body"
      $ checkFun'
        ( nameFromString "<loop body>",
          map (,allowAllAliases (length merge) (length merge)) (staticShapes rettype),
          funParamsToNameInfos mergepat
        )
        (Just consumable)
      $ do
        checkFunParams mergepat
        checkBodyDec $ snd $ bodyDec loopbody

        checkStms (bodyStms loopbody) $ do
          context "In loop body result" $
            checkResult $
              bodyResult loopbody

          context "When matching result of body with loop parameters" $
            matchLoopResult (map fst merge) $
              bodyResult loopbody

          let bound_here =
                [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Scope (Aliases rep) -> [VName]
forall k a. Map k a -> [k]
M.keys (Scope (Aliases rep) -> [VName]) -> Scope (Aliases rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stms (Aliases rep) -> Scope (Aliases rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf (Stms (Aliases rep) -> Scope (Aliases rep))
-> Stms (Aliases rep) -> Scope (Aliases rep)
forall a b. (a -> b) -> a -> b
$ Body (Aliases rep) -> Stms (Aliases rep)
forall rep. Body rep -> Stms rep
bodyStms Body (Aliases rep)
loopbody
          map (`namesSubtract` bound_here)
            <$> mapM (subExpAliasesM . resSubExp) (bodyResult loopbody)
  where
    checkForm :: [Arg] -> LoopForm -> TypeM rep [(VName, Names)]
checkForm [Arg]
mergeargs (ForLoop VName
loopvar IntType
it SubExp
boundexp) = do
      iparam <- VName -> PrimType -> TypeM rep (FParam (Aliases rep))
forall rep.
Checkable rep =>
VName -> PrimType -> TypeM rep (FParam (Aliases rep))
primFParam VName
loopvar (PrimType -> TypeM rep (FParam (Aliases rep)))
-> PrimType -> TypeM rep (FParam (Aliases rep))
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it
      let mergepat = ((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
-> [(Param (FParamInfo rep), SubExp)] -> [Param (FParamInfo rep)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
forall a b. (a, b) -> a
fst [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge
          funparams = Param (FParamInfo rep)
iparam Param (FParamInfo rep)
-> [Param (FParamInfo rep)] -> [Param (FParamInfo rep)]
forall a. a -> [a] -> [a]
: [Param (FParamInfo rep)]
mergepat
          paramts = (Param (FParamInfo rep) -> DeclType)
-> [Param (FParamInfo rep)] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType [Param (FParamInfo rep)]
funparams

      boundarg <- checkArg boundexp
      checkFuncall Nothing paramts $ boundarg : mergeargs
      pure mempty
    checkForm [Arg]
mergeargs (WhileLoop VName
cond) = do
      case ((Param (FParamInfo rep), SubExp) -> Bool)
-> [(Param (FParamInfo rep), SubExp)]
-> Maybe (Param (FParamInfo rep), SubExp)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
cond) (VName -> Bool)
-> ((Param (FParamInfo rep), SubExp) -> VName)
-> (Param (FParamInfo rep), SubExp)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName (Param (FParamInfo rep) -> VName)
-> ((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
-> (Param (FParamInfo rep), SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
forall a b. (a, b) -> a
fst) [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge of
        Just (Param (FParamInfo rep)
condparam, SubExp
_) ->
          Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Param (FParamInfo rep) -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param (FParamInfo rep)
condparam Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Bool) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
            ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
              Text
"Conditional '"
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
cond
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"' of while-loop is not boolean, but "
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Type -> Text
forall a. Pretty a => a -> Text
prettyText (Param (FParamInfo rep) -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param (FParamInfo rep)
condparam)
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"."
        Maybe (Param (FParamInfo rep), SubExp)
Nothing ->
          -- Implies infinite loop, but that's OK.
          () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      let mergepat :: [Param (FParamInfo rep)]
mergepat = ((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
-> [(Param (FParamInfo rep), SubExp)] -> [Param (FParamInfo rep)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
forall a b. (a, b) -> a
fst [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge
          funparams :: [Param (FParamInfo rep)]
funparams = [Param (FParamInfo rep)]
mergepat
          paramts :: [DeclType]
paramts = (Param (FParamInfo rep) -> DeclType)
-> [Param (FParamInfo rep)] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType [Param (FParamInfo rep)]
funparams
      Maybe Name -> [DeclType] -> [Arg] -> TypeM rep ()
forall rep. Maybe Name -> [DeclType] -> [Arg] -> TypeM rep ()
checkFuncall Maybe Name
forall a. Maybe a
Nothing [DeclType]
paramts [Arg]
mergeargs
      [(VName, Names)] -> TypeM rep [(VName, Names)]
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(VName, Names)]
forall a. Monoid a => a
mempty

    checkLoopArgs :: TypeM rep ()
checkLoopArgs = do
      let ([Param (FParamInfo rep)]
params, [SubExp]
args) = [(Param (FParamInfo rep), SubExp)]
-> ([Param (FParamInfo rep)], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge

      argtypes <- (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
args

      let expected = [VName] -> [Param (FParamInfo rep)] -> [SubExp] -> [Type]
forall t. Typed t => [VName] -> [t] -> [SubExp] -> [Type]
expectedTypes ((Param (FParamInfo rep) -> VName)
-> [Param (FParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (FParamInfo rep)]
params) [Param (FParamInfo rep)]
params [SubExp]
args
      unless (expected == argtypes) . bad . TypeError . docText $
        "Loop parameters"
          </> indent 2 (ppTuple' $ map pretty params)
          </> "cannot accept initial values"
          </> indent 2 (ppTuple' $ map pretty args)
          </> "of types"
          </> indent 2 (ppTuple' $ map pretty argtypes)
checkExp (WithAcc [WithAccInput (Aliases rep)]
inputs Lambda (Aliases rep)
lam) = do
  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Param (LParamInfo rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Lambda (Aliases rep) -> [LParam (Aliases rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda (Aliases rep)
lam) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
num_accs) (TypeM rep () -> TypeM rep ())
-> (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    Int -> Text
forall a. Pretty a => a -> Text
prettyText ([Param (LParamInfo rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Lambda (Aliases rep) -> [LParam (Aliases rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda (Aliases rep)
lam))
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" parameters, but "
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
num_accs
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" accumulators."

  let cert_params :: [Param (LParamInfo rep)]
cert_params = Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a. Int -> [a] -> [a]
take Int
num_accs ([Param (LParamInfo rep)] -> [Param (LParamInfo rep)])
-> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
forall a b. (a -> b) -> a -> b
$ Lambda (Aliases rep) -> [LParam (Aliases rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda (Aliases rep)
lam
  acc_args <- [(WithAccInput (Aliases rep), Param (LParamInfo rep))]
-> ((WithAccInput (Aliases rep), Param (LParamInfo rep))
    -> TypeM rep Arg)
-> TypeM rep [Arg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([WithAccInput (Aliases rep)]
-> [Param (LParamInfo rep)]
-> [(WithAccInput (Aliases rep), Param (LParamInfo rep))]
forall a b. [a] -> [b] -> [(a, b)]
zip [WithAccInput (Aliases rep)]
inputs [Param (LParamInfo rep)]
cert_params) (((WithAccInput (Aliases rep), Param (LParamInfo rep))
  -> TypeM rep Arg)
 -> TypeM rep [Arg])
-> ((WithAccInput (Aliases rep), Param (LParamInfo rep))
    -> TypeM rep Arg)
-> TypeM rep [Arg]
forall a b. (a -> b) -> a -> b
$ \((Shape
shape, [VName]
arrs, Maybe (Lambda (Aliases rep), [SubExp])
op), Param (LParamInfo rep)
p) -> do
    (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64]) (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape)
    elem_ts <- [VName] -> (VName -> TypeM rep Type) -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> TypeM rep Type) -> TypeM rep [Type])
-> (VName -> TypeM rep Type) -> TypeM rep [Type]
forall a b. (a -> b) -> a -> b
$ \VName
arr -> do
      arr_t <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
      unless (shapeDims shape `isPrefixOf` arrayDims arr_t) $
        bad . TypeError $
          prettyText arr <> " is not an array of outer shape " <> prettyText shape
      consume =<< lookupAliases arr
      pure $ stripArray (shapeRank shape) arr_t

    case op of
      Just (Lambda (Aliases rep)
op_lam, [SubExp]
nes) -> do
        let mkArrArg :: a -> (a, b)
mkArrArg a
t = (a
t, b
forall a. Monoid a => a
mempty)
        nes_ts <- (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp [SubExp]
nes
        unless (nes_ts == lambdaReturnType op_lam) $
          bad . TypeError . T.unlines $
            [ "Accumulator operator return type: " <> prettyText (lambdaReturnType op_lam),
              "Type of neutral elements: " <> prettyText nes_ts
            ]
        checkLambda op_lam $
          replicate (shapeRank shape) (Prim int64, mempty)
            ++ map mkArrArg (elem_ts ++ elem_ts)
      Maybe (Lambda (Aliases rep), [SubExp])
Nothing ->
        () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    pure (Acc (paramName p) shape elem_ts NoUniqueness, mempty)

  checkAnyLambda False lam $ replicate num_accs (Prim Unit, mempty) ++ acc_args
  where
    num_accs :: Int
num_accs = [WithAccInput (Aliases rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput (Aliases rep)]
inputs
checkExp (Op Op (Aliases rep)
op) = do
  checker <- (Env rep -> OpC rep (Aliases rep) -> TypeM rep ())
-> TypeM rep (OpC rep (Aliases rep) -> TypeM rep ())
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env rep -> OpC rep (Aliases rep) -> TypeM rep ()
Env rep -> Op (Aliases rep) -> TypeM rep ()
forall rep. Env rep -> Op (Aliases rep) -> TypeM rep ()
envCheckOp
  checker op

checkSOACArrayArgs ::
  (Checkable rep) =>
  SubExp ->
  [VName] ->
  TypeM rep [Arg]
checkSOACArrayArgs :: forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
checkSOACArrayArgs SubExp
width = (VName -> TypeM rep Arg) -> [VName] -> TypeM rep [Arg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> TypeM rep Arg
checkSOACArrayArg
  where
    checkSOACArrayArg :: VName -> TypeM rep Arg
checkSOACArrayArg VName
v = do
      (t, als) <- SubExp -> TypeM rep Arg
forall rep. Checkable rep => SubExp -> TypeM rep Arg
checkArg (SubExp -> TypeM rep Arg) -> SubExp -> TypeM rep Arg
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
      case t of
        Acc {} -> Arg -> TypeM rep Arg
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
t, Names
als)
        Array {} -> do
          let argSize :: SubExp
argSize = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 Type
t
          Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
argSize SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
width) (TypeM rep () -> TypeM rep ())
-> (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
            Text
"SOAC argument "
              Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v
              Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" has outer size "
              Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
argSize
              Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", but width of SOAC is "
              Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
width
          Arg -> TypeM rep Arg
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType Type
t, Names
als)
        Type
_ ->
          ErrorCase rep -> TypeM rep Arg
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep Arg)
-> (Text -> ErrorCase rep) -> Text -> TypeM rep Arg
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep Arg) -> Text -> TypeM rep Arg
forall a b. (a -> b) -> a -> b
$
            Text
"SOAC argument " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" is not an array"

checkType ::
  (Checkable rep) =>
  TypeBase Shape u ->
  TypeM rep ()
checkType :: forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
checkType (Mem (ScalarSpace [SubExp]
d PrimType
_)) = (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64]) [SubExp]
d
checkType (Acc VName
cert Shape
shape [Type]
ts u
_) = do
  [Type] -> VName -> TypeM rep ()
forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
requireI [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Unit] VName
cert
  (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64]) ([SubExp] -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
  (Type -> TypeM rep ()) -> [Type] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Type -> TypeM rep ()
forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
checkType [Type]
ts
checkType TypeBase Shape u
t = (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp ([SubExp] -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ TypeBase Shape u -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims TypeBase Shape u
t

checkExtType ::
  (Checkable rep) =>
  TypeBase ExtShape u ->
  TypeM rep ()
checkExtType :: forall rep u. Checkable rep => TypeBase ExtShape u -> TypeM rep ()
checkExtType = (Ext SubExp -> TypeM rep ()) -> [Ext SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ext SubExp -> TypeM rep ()
forall {rep}. Checkable rep => Ext SubExp -> TypeM rep ()
checkExtDim ([Ext SubExp] -> TypeM rep ())
-> (TypeBase ExtShape u -> [Ext SubExp])
-> TypeBase ExtShape u
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtShape -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ExtShape -> [Ext SubExp])
-> (TypeBase ExtShape u -> ExtShape)
-> TypeBase ExtShape u
-> [Ext SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape u -> ExtShape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape
  where
    checkExtDim :: Ext SubExp -> TypeM rep ()
checkExtDim (Free SubExp
se) = TypeM rep Type -> TypeM rep ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (TypeM rep Type -> TypeM rep ()) -> TypeM rep Type -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
se
    checkExtDim (Ext Int
_) = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

checkCmpOp ::
  (Checkable rep) =>
  CmpOp ->
  SubExp ->
  SubExp ->
  TypeM rep ()
checkCmpOp :: forall rep.
Checkable rep =>
CmpOp -> SubExp -> SubExp -> TypeM rep ()
checkCmpOp (CmpEq PrimType
t) SubExp
x SubExp
y = do
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t] SubExp
x
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t] SubExp
y
checkCmpOp (CmpUlt IntType
t) SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (IntType -> PrimType
IntType IntType
t) SubExp
x SubExp
y
checkCmpOp (CmpUle IntType
t) SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (IntType -> PrimType
IntType IntType
t) SubExp
x SubExp
y
checkCmpOp (CmpSlt IntType
t) SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (IntType -> PrimType
IntType IntType
t) SubExp
x SubExp
y
checkCmpOp (CmpSle IntType
t) SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (IntType -> PrimType
IntType IntType
t) SubExp
x SubExp
y
checkCmpOp (FCmpLt FloatType
t) SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (FloatType -> PrimType
FloatType FloatType
t) SubExp
x SubExp
y
checkCmpOp (FCmpLe FloatType
t) SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs (FloatType -> PrimType
FloatType FloatType
t) SubExp
x SubExp
y
checkCmpOp CmpOp
CmpLlt SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs PrimType
Bool SubExp
x SubExp
y
checkCmpOp CmpOp
CmpLle SubExp
x SubExp
y = PrimType -> SubExp -> SubExp -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs PrimType
Bool SubExp
x SubExp
y

checkBinOpArgs ::
  (Checkable rep) =>
  PrimType ->
  SubExp ->
  SubExp ->
  TypeM rep ()
checkBinOpArgs :: forall rep.
Checkable rep =>
PrimType -> SubExp -> SubExp -> TypeM rep ()
checkBinOpArgs PrimType
t SubExp
e1 SubExp
e2 = do
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t] SubExp
e1
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t] SubExp
e2

checkPatElem ::
  (Checkable rep) =>
  PatElem (LetDec rep) ->
  TypeM rep ()
checkPatElem :: forall rep. Checkable rep => PatElem (LetDec rep) -> TypeM rep ()
checkPatElem (PatElem VName
name LetDec rep
dec) =
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"When checking pattern element " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
name) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    VName -> LetDec rep -> TypeM rep ()
forall rep. Checkable rep => VName -> LetDec rep -> TypeM rep ()
checkLetBoundDec VName
name LetDec rep
dec

checkFlatDimIndex ::
  (Checkable rep) =>
  FlatDimIndex SubExp ->
  TypeM rep ()
checkFlatDimIndex :: forall rep. Checkable rep => FlatDimIndex SubExp -> TypeM rep ()
checkFlatDimIndex (FlatDimIndex SubExp
n SubExp
s) = (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64]) [SubExp
n, SubExp
s]

checkFlatSlice ::
  (Checkable rep) =>
  FlatSlice SubExp ->
  TypeM rep ()
checkFlatSlice :: forall rep. Checkable rep => FlatSlice SubExp -> TypeM rep ()
checkFlatSlice (FlatSlice SubExp
offset [FlatDimIndex SubExp]
idxs) = do
  [Type] -> SubExp -> TypeM rep ()
forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
require [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64] SubExp
offset
  (FlatDimIndex SubExp -> TypeM rep ())
-> [FlatDimIndex SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ FlatDimIndex SubExp -> TypeM rep ()
forall rep. Checkable rep => FlatDimIndex SubExp -> TypeM rep ()
checkFlatDimIndex [FlatDimIndex SubExp]
idxs

checkStm ::
  (Checkable rep) =>
  Stm (Aliases rep) ->
  TypeM rep a ->
  TypeM rep a
checkStm :: forall rep a.
Checkable rep =>
Stm (Aliases rep) -> TypeM rep a -> TypeM rep a
checkStm stm :: Stm (Aliases rep)
stm@(Let Pat (LetDec (Aliases rep))
pat StmAux (ExpDec (Aliases rep))
aux Exp (Aliases rep)
e) TypeM rep a
m = do
  let Certs [VName]
cs = StmAux (VarAliases, ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux (VarAliases, ExpDec rep)
StmAux (ExpDec (Aliases rep))
aux
      (VarAliases
_, ExpDec rep
dec) = StmAux (VarAliases, ExpDec rep) -> (VarAliases, ExpDec rep)
forall dec. StmAux dec -> dec
stmAuxDec StmAux (VarAliases, ExpDec rep)
StmAux (ExpDec (Aliases rep))
aux
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context Text
"When checking certificates" (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ (VName -> TypeM rep ()) -> [VName] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Type] -> VName -> TypeM rep ()
forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
requireI [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Unit]) [VName]
cs
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context Text
"When checking expression annotation" (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ ExpDec rep -> TypeM rep ()
forall rep. Checkable rep => ExpDec rep -> TypeM rep ()
checkExpDec ExpDec rep
dec
  Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"When matching\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Pat (VarAliases, LetDec rep) -> Text
forall a. Pretty a => Text -> a -> Text
message Text
"  " Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\nwith\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Exp (Aliases rep) -> Text
forall a. Pretty a => Text -> a -> Text
message Text
"  " Exp (Aliases rep)
e) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
forall rep.
Checkable rep =>
Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
matchPat Pat (LetDec (Aliases rep))
pat Exp (Aliases rep)
e
  Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
binding (Stm (Aliases rep) -> Scope (Aliases rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm (Aliases rep)
stm) (TypeM rep a -> TypeM rep a) -> TypeM rep a -> TypeM rep a
forall a b. (a -> b) -> a -> b
$ do
    (PatElem (LetDec rep) -> TypeM rep ())
-> [PatElem (LetDec rep)] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ PatElem (LetDec rep) -> TypeM rep ()
forall rep. Checkable rep => PatElem (LetDec rep) -> TypeM rep ()
checkPatElem (Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems (Pat (LetDec rep) -> [PatElem (LetDec rep)])
-> Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall a b. (a -> b) -> a -> b
$ Pat (VarAliases, LetDec rep) -> Pat (LetDec rep)
forall a. Pat (VarAliases, a) -> Pat a
removePatAliases Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat)
    TypeM rep a
m

matchExtPat ::
  (Checkable rep) =>
  Pat (LetDec (Aliases rep)) ->
  [ExtType] ->
  TypeM rep ()
matchExtPat :: forall rep.
Checkable rep =>
Pat (LetDec (Aliases rep)) -> [ExtType] -> TypeM rep ()
matchExtPat Pat (LetDec (Aliases rep))
pat [ExtType]
ts =
  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Pat (VarAliases, LetDec rep) -> [ExtType]
forall dec. Typed dec => Pat dec -> [ExtType]
expExtTypesFromPat Pat (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat [ExtType] -> [ExtType] -> Bool
forall a. Eq a => a -> a -> Bool
== [ExtType]
ts) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
      Pat (LetDec (Aliases rep))
-> [ExtType] -> Maybe String -> ErrorCase rep
forall rep.
Pat (LetDec (Aliases rep))
-> [ExtType] -> Maybe String -> ErrorCase rep
InvalidPatError Pat (LetDec (Aliases rep))
pat [ExtType]
ts Maybe String
forall a. Maybe a
Nothing

matchExtReturnType ::
  (Checkable rep) =>
  [ExtType] ->
  Result ->
  TypeM rep ()
matchExtReturnType :: forall rep. Checkable rep => [ExtType] -> Result -> TypeM rep ()
matchExtReturnType [ExtType]
rettype Result
res = do
  ts <- (SubExpRes -> TypeM rep Type) -> Result -> TypeM rep [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExpRes -> TypeM rep Type
forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
subExpResType Result
res
  matchExtReturns rettype res ts

matchExtBranchType ::
  (Checkable rep) =>
  [ExtType] ->
  Body (Aliases rep) ->
  TypeM rep ()
matchExtBranchType :: forall rep.
Checkable rep =>
[ExtType] -> Body (Aliases rep) -> TypeM rep ()
matchExtBranchType [ExtType]
rettype (Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
res) = do
  ts <- ExtendedScope (Aliases rep) (TypeM rep) [Type]
-> Scope (Aliases rep) -> TypeM rep [Type]
forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope ((SubExpRes -> ExtendedScope (Aliases rep) (TypeM rep) Type)
-> Result -> ExtendedScope (Aliases rep) (TypeM rep) [Type]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse SubExpRes -> ExtendedScope (Aliases rep) (TypeM rep) Type
forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
subExpResType Result
res) Scope (Aliases rep)
stmscope
  matchExtReturns rettype res ts
  where
    stmscope :: Scope (Aliases rep)
stmscope = Stms (Aliases rep) -> Scope (Aliases rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms (Aliases rep)
stms

matchExtReturns :: [ExtType] -> Result -> [Type] -> TypeM rep ()
matchExtReturns :: forall rep. [ExtType] -> Result -> [Type] -> TypeM rep ()
matchExtReturns [ExtType]
rettype Result
res [Type]
ts = do
  let problem :: TypeM rep a
      problem :: forall rep a. TypeM rep a
problem =
        ErrorCase rep -> TypeM rep a
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep a)
-> ([Text] -> ErrorCase rep) -> [Text] -> TypeM rep a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> ErrorCase rep)
-> ([Text] -> Text) -> [Text] -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Text
T.unlines ([Text] -> TypeM rep a) -> [Text] -> TypeM rep a
forall a b. (a -> b) -> a -> b
$
          [ Text
"Type annotation is",
            Text
"  " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [ExtType] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [ExtType]
rettype,
            Text
"But result returns type",
            Text
"  " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Type] -> Text
forall a. Pretty a => [a] -> Text
prettyTuple [Type]
ts
          ]

  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Result -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
res Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [ExtType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
rettype) TypeM rep ()
forall rep a. TypeM rep a
problem

  let ctx_vals :: [(SubExpRes, Type)]
ctx_vals = Result -> [Type] -> [(SubExpRes, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
res [Type]
ts
      instantiateExt :: Int -> TypeM rep SubExp
instantiateExt Int
i = case Int -> [(SubExpRes, Type)] -> Maybe (SubExpRes, Type)
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth Int
i [(SubExpRes, Type)]
ctx_vals of
        Just (SubExpRes Certs
_ SubExp
se, Prim (IntType IntType
Int64)) -> SubExp -> TypeM rep SubExp
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
        Maybe (SubExpRes, Type)
_ -> TypeM rep SubExp
forall rep a. TypeM rep a
problem

  rettype' <- (Int -> TypeM rep SubExp) -> [ExtType] -> TypeM rep [Type]
forall (m :: * -> *) u.
Monad m =>
(Int -> m SubExp) -> [TypeBase ExtShape u] -> m [TypeBase Shape u]
instantiateShapes Int -> TypeM rep SubExp
instantiateExt [ExtType]
rettype

  unless (rettype' == ts) problem

validApply ::
  (ArrayShape shape) =>
  [TypeBase shape Uniqueness] ->
  [TypeBase shape NoUniqueness] ->
  Bool
validApply :: forall shape.
ArrayShape shape =>
[TypeBase shape Uniqueness]
-> [TypeBase shape NoUniqueness] -> Bool
validApply [TypeBase shape Uniqueness]
expected [TypeBase shape NoUniqueness]
got =
  [TypeBase shape NoUniqueness] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase shape NoUniqueness]
got Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [TypeBase shape Uniqueness] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase shape Uniqueness]
expected
    Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and
      ( (TypeBase Rank NoUniqueness -> TypeBase Rank NoUniqueness -> Bool)
-> [TypeBase Rank NoUniqueness]
-> [TypeBase Rank NoUniqueness]
-> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          TypeBase Rank NoUniqueness -> TypeBase Rank NoUniqueness -> Bool
forall u shape.
(Ord u, ArrayShape shape) =>
TypeBase shape u -> TypeBase shape u -> Bool
subtypeOf
          ((TypeBase shape NoUniqueness -> TypeBase Rank NoUniqueness)
-> [TypeBase shape NoUniqueness] -> [TypeBase Rank NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase shape NoUniqueness -> TypeBase Rank NoUniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
rankShaped [TypeBase shape NoUniqueness]
got)
          ((TypeBase shape Uniqueness -> TypeBase Rank NoUniqueness)
-> [TypeBase shape Uniqueness] -> [TypeBase Rank NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Rank Uniqueness -> TypeBase Rank NoUniqueness
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl (TypeBase Rank Uniqueness -> TypeBase Rank NoUniqueness)
-> (TypeBase shape Uniqueness -> TypeBase Rank Uniqueness)
-> TypeBase shape Uniqueness
-> TypeBase Rank NoUniqueness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase shape Uniqueness -> TypeBase Rank Uniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
rankShaped) [TypeBase shape Uniqueness]
expected)
      )

type Arg = (Type, Names)

argType :: Arg -> Type
argType :: Arg -> Type
argType (Type
t, Names
_) = Type
t

-- | Remove all aliases from the 'Arg'.
argAliases :: Arg -> Names
argAliases :: Arg -> Names
argAliases (Type
_, Names
als) = Names
als

noArgAliases :: Arg -> Arg
noArgAliases :: Arg -> Arg
noArgAliases (Type
t, Names
_) = (Type
t, Names
forall a. Monoid a => a
mempty)

checkArg ::
  (Checkable rep) =>
  SubExp ->
  TypeM rep Arg
checkArg :: forall rep. Checkable rep => SubExp -> TypeM rep Arg
checkArg SubExp
arg = do
  argt <- SubExp -> TypeM rep Type
forall rep. Checkable rep => SubExp -> TypeM rep Type
checkSubExp SubExp
arg
  als <- subExpAliasesM arg
  pure (argt, als)

checkFuncall ::
  Maybe Name ->
  [DeclType] ->
  [Arg] ->
  TypeM rep ()
checkFuncall :: forall rep. Maybe Name -> [DeclType] -> [Arg] -> TypeM rep ()
checkFuncall Maybe Name
fname [DeclType]
paramts [Arg]
args = do
  let argts :: [Type]
argts = (Arg -> Type) -> [Arg] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> Type
argType [Arg]
args
  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([DeclType] -> [Type] -> Bool
forall shape.
ArrayShape shape =>
[TypeBase shape Uniqueness]
-> [TypeBase shape NoUniqueness] -> Bool
validApply [DeclType]
paramts [Type]
argts) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
      Maybe Name -> [Type] -> [Type] -> ErrorCase rep
forall rep. Maybe Name -> [Type] -> [Type] -> ErrorCase rep
ParameterMismatch Maybe Name
fname ((DeclType -> Type) -> [DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Type
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl [DeclType]
paramts) ([Type] -> ErrorCase rep) -> [Type] -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
        (Arg -> Type) -> [Arg] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> Type
argType [Arg]
args
  [DeclType] -> [Arg] -> TypeM rep ()
forall rep. [DeclType] -> [Arg] -> TypeM rep ()
consumeArgs [DeclType]
paramts [Arg]
args

consumeArgs ::
  [DeclType] ->
  [Arg] ->
  TypeM rep ()
consumeArgs :: forall rep. [DeclType] -> [Arg] -> TypeM rep ()
consumeArgs [DeclType]
paramts [Arg]
args =
  [(Diet, Arg)] -> ((Diet, Arg) -> TypeM rep ()) -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Diet] -> [Arg] -> [(Diet, Arg)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((DeclType -> Diet) -> [DeclType] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
diet [DeclType]
paramts) [Arg]
args) (((Diet, Arg) -> TypeM rep ()) -> TypeM rep ())
-> ((Diet, Arg) -> TypeM rep ()) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ \(Diet
d, (Type
_, Names
als)) ->
    [Occurence] -> TypeM rep ()
forall rep. [Occurence] -> TypeM rep ()
occur [Names -> Occurence
consumption (Names -> Diet -> Names
forall {p}. Monoid p => p -> Diet -> p
consumeArg Names
als Diet
d)]
  where
    consumeArg :: p -> Diet -> p
consumeArg p
als Diet
Consume = p
als
    consumeArg p
_ Diet
_ = p
forall a. Monoid a => a
mempty

-- The boolean indicates whether we only allow consumption of
-- parameters.
checkAnyLambda ::
  (Checkable rep) => Bool -> Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
checkAnyLambda :: forall rep.
Checkable rep =>
Bool -> Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
checkAnyLambda Bool
soac (Lambda [LParam (Aliases rep)]
params [Type]
rettype Body (Aliases rep)
body) [Arg]
args = do
  let fname :: Name
fname = String -> Name
nameFromString String
"<anonymous>"
  if [Param (LParamInfo rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Arg] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Arg]
args
    then do
      -- Consumption for this is done explicitly elsewhere.
      Maybe Name -> [DeclType] -> [Arg] -> TypeM rep ()
forall rep. Maybe Name -> [DeclType] -> [Arg] -> TypeM rep ()
checkFuncall
        Maybe Name
forall a. Maybe a
Nothing
        ((Param (LParamInfo rep) -> DeclType)
-> [Param (LParamInfo rep)] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Type -> Uniqueness -> DeclType
forall shape.
TypeBase shape NoUniqueness
-> Uniqueness -> TypeBase shape Uniqueness
`toDecl` Uniqueness
Nonunique) (Type -> DeclType)
-> (Param (LParamInfo rep) -> Type)
-> Param (LParamInfo rep)
-> DeclType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (LParamInfo rep) -> Type
forall dec. Typed dec => Param dec -> Type
paramType) [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params)
        ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ (Arg -> Arg) -> [Arg] -> [Arg]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> Arg
noArgAliases [Arg]
args
      let consumable :: Maybe [(VName, Names)]
consumable =
            if Bool
soac
              then [(VName, Names)] -> Maybe [(VName, Names)]
forall a. a -> Maybe a
Just ([(VName, Names)] -> Maybe [(VName, Names)])
-> [(VName, Names)] -> Maybe [(VName, Names)]
forall a b. (a -> b) -> a -> b
$ [VName] -> [Names] -> [(VName, Names)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params) ((Arg -> Names) -> [Arg] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> Names
argAliases [Arg]
args)
              else Maybe [(VName, Names)]
forall a. Maybe a
Nothing
          params' :: [(VName, NameInfo (Aliases rep))]
params' =
            [(Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
param, LParamInfo (Aliases rep) -> NameInfo (Aliases rep)
forall rep. LParamInfo rep -> NameInfo rep
LParamName (LParamInfo (Aliases rep) -> NameInfo (Aliases rep))
-> LParamInfo (Aliases rep) -> NameInfo (Aliases rep)
forall a b. (a -> b) -> a -> b
$ Param (LParamInfo rep) -> LParamInfo rep
forall dec. Param dec -> dec
paramDec Param (LParamInfo rep)
param) | Param (LParamInfo rep)
param <- [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params]
      Name -> [VName] -> TypeM rep ()
forall rep. Name -> [VName] -> TypeM rep ()
checkNoDuplicateParams Name
fname ([VName] -> TypeM rep ()) -> [VName] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ (Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params
      Scope (Aliases rep) -> TypeM rep () -> TypeM rep ()
forall rep a.
Checkable rep =>
Scope (Aliases rep) -> TypeM rep a -> TypeM rep a
binding ([(VName, NameInfo (Aliases rep))] -> Scope (Aliases rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, NameInfo (Aliases rep))]
params') (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        (TypeM rep () -> TypeM rep ())
-> ([(VName, Names)] -> TypeM rep () -> TypeM rep ())
-> Maybe [(VName, Names)]
-> TypeM rep ()
-> TypeM rep ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeM rep () -> TypeM rep ()
forall a. a -> a
id [(VName, Names)] -> TypeM rep () -> TypeM rep ()
forall rep a. [(VName, Names)] -> TypeM rep a -> TypeM rep a
consumeOnlyParams Maybe [(VName, Names)]
consumable (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
          [Param (LParamInfo rep)] -> TypeM rep ()
forall rep. Checkable rep => [LParam rep] -> TypeM rep ()
checkLambdaParams [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params
          (Type -> TypeM rep ()) -> [Type] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Type -> TypeM rep ()
forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
checkType [Type]
rettype
          [Type] -> Body (Aliases rep) -> TypeM rep ()
forall rep.
Checkable rep =>
[Type] -> Body (Aliases rep) -> TypeM rep ()
checkLambdaBody [Type]
rettype Body (Aliases rep)
body
    else
      ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Text
"Anonymous function defined with "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText ([Param (LParamInfo rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" parameters:\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Param (LParamInfo rep)] -> Text
forall a. Pretty a => a -> Text
prettyText [Param (LParamInfo rep)]
[LParam (Aliases rep)]
params
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\nbut expected to take "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText ([Arg] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Arg]
args)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" arguments."

checkLambda :: (Checkable rep) => Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
checkLambda :: forall rep.
Checkable rep =>
Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
checkLambda = Bool -> Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
forall rep.
Checkable rep =>
Bool -> Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
checkAnyLambda Bool
True

checkPrimExp :: (Checkable rep) => PrimExp VName -> TypeM rep ()
checkPrimExp :: forall rep. Checkable rep => PrimExp VName -> TypeM rep ()
checkPrimExp ValueExp {} = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkPrimExp (LeafExp VName
v PrimType
pt) = [Type] -> VName -> TypeM rep ()
forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
requireI [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
pt] VName
v
checkPrimExp (BinOpExp BinOp
op PrimExp VName
x PrimExp VName
y) = do
  PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp (BinOp -> PrimType
binOpType BinOp
op) PrimExp VName
x
  PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp (BinOp -> PrimType
binOpType BinOp
op) PrimExp VName
y
checkPrimExp (CmpOpExp CmpOp
op PrimExp VName
x PrimExp VName
y) = do
  PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp (CmpOp -> PrimType
cmpOpType CmpOp
op) PrimExp VName
x
  PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp (CmpOp -> PrimType
cmpOpType CmpOp
op) PrimExp VName
y
checkPrimExp (UnOpExp UnOp
op PrimExp VName
x) = PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp (UnOp -> PrimType
unOpType UnOp
op) PrimExp VName
x
checkPrimExp (ConvOpExp ConvOp
op PrimExp VName
x) = PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp ((PrimType, PrimType) -> PrimType
forall a b. (a, b) -> a
fst ((PrimType, PrimType) -> PrimType)
-> (PrimType, PrimType) -> PrimType
forall a b. (a -> b) -> a -> b
$ ConvOp -> (PrimType, PrimType)
convOpType ConvOp
op) PrimExp VName
x
checkPrimExp (FunExp Text
h [PrimExp VName]
args PrimType
t) = do
  (h_ts, h_ret, _) <-
    TypeM rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> (([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
    -> TypeM
         rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
-> Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> TypeM rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
      (ErrorCase rep
-> TypeM rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep
 -> TypeM
      rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
-> ErrorCase rep
-> TypeM rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall a b. (a -> b) -> a -> b
$ Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> ErrorCase rep) -> Text -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$ Text
"Unknown function: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
h)
      ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> TypeM rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      (Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
 -> TypeM
      rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
-> Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> TypeM rep ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall a b. (a -> b) -> a -> b
$ Text
-> Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
h Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
primFuns
  when (length h_ts /= length args) . bad . TypeError $
    "Function expects "
      <> prettyText (length h_ts)
      <> " parameters, but given "
      <> prettyText (length args)
      <> " arguments."
  when (h_ret /= t) . bad . TypeError $
    "Function return annotation is "
      <> prettyText t
      <> ", but expected "
      <> prettyText h_ret
  zipWithM_ requirePrimExp h_ts args

requirePrimExp :: (Checkable rep) => PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp :: forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
requirePrimExp PrimType
t PrimExp VName
e = Text -> TypeM rep () -> TypeM rep ()
forall rep a. Text -> TypeM rep a -> TypeM rep a
context (Text
"in PrimExp " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> PrimExp VName -> Text
forall a. Pretty a => a -> Text
prettyText PrimExp VName
e) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
  PrimExp VName -> TypeM rep ()
forall rep. Checkable rep => PrimExp VName -> TypeM rep ()
checkPrimExp PrimExp VName
e
  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (PrimExp VName -> PrimType
forall v. PrimExp v -> PrimType
primExpType PrimExp VName
e PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
t) (TypeM rep () -> TypeM rep ())
-> (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    PrimExp VName -> Text
forall a. Pretty a => a -> Text
prettyText PrimExp VName
e Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" must have type " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> PrimType -> Text
forall a. Pretty a => a -> Text
prettyText PrimType
t

-- | The class of representations that can be type-checked.
class (AliasableRep rep, TypedOp (OpC rep)) => Checkable rep where
  checkExpDec :: ExpDec rep -> TypeM rep ()
  checkBodyDec :: BodyDec rep -> TypeM rep ()
  checkFParamDec :: VName -> FParamInfo rep -> TypeM rep ()
  checkLParamDec :: VName -> LParamInfo rep -> TypeM rep ()
  checkLetBoundDec :: VName -> LetDec rep -> TypeM rep ()
  checkRetType :: [RetType rep] -> TypeM rep ()
  matchPat :: Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
  primFParam :: VName -> PrimType -> TypeM rep (FParam (Aliases rep))
  matchReturnType :: [RetType rep] -> Result -> TypeM rep ()
  matchBranchType :: [BranchType rep] -> Body (Aliases rep) -> TypeM rep ()
  matchLoopResult :: [FParam (Aliases rep)] -> Result -> TypeM rep ()

  -- | Used at top level; can be locally changed with 'checkOpWith'.
  checkOp :: Op (Aliases rep) -> TypeM rep ()

  default checkExpDec :: (ExpDec rep ~ ()) => ExpDec rep -> TypeM rep ()
  checkExpDec = () -> TypeM rep ()
ExpDec rep -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

  default checkBodyDec :: (BodyDec rep ~ ()) => BodyDec rep -> TypeM rep ()
  checkBodyDec = () -> TypeM rep ()
BodyDec rep -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

  default checkFParamDec :: (FParamInfo rep ~ DeclType) => VName -> FParamInfo rep -> TypeM rep ()
  checkFParamDec VName
_ = DeclType -> TypeM rep ()
FParamInfo rep -> TypeM rep ()
forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
checkType

  default checkLParamDec :: (LParamInfo rep ~ Type) => VName -> LParamInfo rep -> TypeM rep ()
  checkLParamDec VName
_ = Type -> TypeM rep ()
LParamInfo rep -> TypeM rep ()
forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
checkType

  default checkLetBoundDec :: (LetDec rep ~ Type) => VName -> LetDec rep -> TypeM rep ()
  checkLetBoundDec VName
_ = Type -> TypeM rep ()
LetDec rep -> TypeM rep ()
forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
checkType

  default checkRetType :: (RetType rep ~ DeclExtType) => [RetType rep] -> TypeM rep ()
  checkRetType = (RetType rep -> TypeM rep ()) -> [RetType rep] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((RetType rep -> TypeM rep ()) -> [RetType rep] -> TypeM rep ())
-> (RetType rep -> TypeM rep ()) -> [RetType rep] -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ DeclExtType -> TypeM rep ()
forall rep u. Checkable rep => TypeBase ExtShape u -> TypeM rep ()
checkExtType (DeclExtType -> TypeM rep ())
-> (DeclExtType -> DeclExtType) -> DeclExtType -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DeclExtType -> DeclExtType
forall t. DeclExtTyped t => t -> DeclExtType
declExtTypeOf

  default matchPat :: Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
  matchPat Pat (LetDec (Aliases rep))
pat = Pat (LetDec (Aliases rep)) -> [ExtType] -> TypeM rep ()
forall rep.
Checkable rep =>
Pat (LetDec (Aliases rep)) -> [ExtType] -> TypeM rep ()
matchExtPat Pat (LetDec (Aliases rep))
pat ([ExtType] -> TypeM rep ())
-> (Exp (Aliases rep) -> TypeM rep [ExtType])
-> Exp (Aliases rep)
-> TypeM rep ()
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Exp (Aliases rep) -> TypeM rep [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [ExtType]
expExtType

  default primFParam :: (FParamInfo rep ~ DeclType) => VName -> PrimType -> TypeM rep (FParam (Aliases rep))
  primFParam VName
name PrimType
t = FParam (Aliases rep) -> TypeM rep (FParam (Aliases rep))
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FParam (Aliases rep) -> TypeM rep (FParam (Aliases rep)))
-> FParam (Aliases rep) -> TypeM rep (FParam (Aliases rep))
forall a b. (a -> b) -> a -> b
$ Attrs -> VName -> DeclType -> Param DeclType
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
name (PrimType -> DeclType
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t)

  default matchReturnType :: (RetType rep ~ DeclExtType) => [RetType rep] -> Result -> TypeM rep ()
  matchReturnType = [ExtType] -> Result -> TypeM rep ()
forall rep. Checkable rep => [ExtType] -> Result -> TypeM rep ()
matchExtReturnType ([ExtType] -> Result -> TypeM rep ())
-> ([DeclExtType] -> [ExtType])
-> [DeclExtType]
-> Result
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl

  default matchBranchType :: (BranchType rep ~ ExtType) => [BranchType rep] -> Body (Aliases rep) -> TypeM rep ()
  matchBranchType = [ExtType] -> Body (Aliases rep) -> TypeM rep ()
[BranchType rep] -> Body (Aliases rep) -> TypeM rep ()
forall rep.
Checkable rep =>
[ExtType] -> Body (Aliases rep) -> TypeM rep ()
matchExtBranchType

  default matchLoopResult ::
    (FParamInfo rep ~ DeclType) =>
    [FParam (Aliases rep)] ->
    Result ->
    TypeM rep ()
  matchLoopResult = [Param DeclType] -> Result -> TypeM rep ()
[FParam (Aliases rep)] -> Result -> TypeM rep ()
forall rep.
Checkable rep =>
[Param DeclType] -> Result -> TypeM rep ()
matchLoopResultExt