-- | Check that a value definition does not violate any consumption
-- constraints.
module Language.Futhark.TypeChecker.Consumption
  ( checkValDef,
  )
where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifoldable
import Data.Bifunctor
import Data.DList qualified as DL
import Data.Foldable
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.Util.Pretty hiding (space)
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Monad (Notes, TypeError (..), withIndexLink)
import Prelude hiding (mod)

type Names = S.Set VName

-- | A variable that is aliased.  Can be still in-scope, or have gone
-- out of scope and be free.  In the latter case, it behaves more like
-- an equivalence class.  See uniqueness-error18.fut for an example of
-- why this is necessary.
data Alias
  = AliasBound {Alias -> VName
aliasVar :: VName}
  | AliasFree {aliasVar :: VName}
  deriving (Alias -> Alias -> Bool
(Alias -> Alias -> Bool) -> (Alias -> Alias -> Bool) -> Eq Alias
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Alias -> Alias -> Bool
== :: Alias -> Alias -> Bool
$c/= :: Alias -> Alias -> Bool
/= :: Alias -> Alias -> Bool
Eq, Eq Alias
Eq Alias =>
(Alias -> Alias -> Ordering)
-> (Alias -> Alias -> Bool)
-> (Alias -> Alias -> Bool)
-> (Alias -> Alias -> Bool)
-> (Alias -> Alias -> Bool)
-> (Alias -> Alias -> Alias)
-> (Alias -> Alias -> Alias)
-> Ord Alias
Alias -> Alias -> Bool
Alias -> Alias -> Ordering
Alias -> Alias -> Alias
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 :: Alias -> Alias -> Ordering
compare :: Alias -> Alias -> Ordering
$c< :: Alias -> Alias -> Bool
< :: Alias -> Alias -> Bool
$c<= :: Alias -> Alias -> Bool
<= :: Alias -> Alias -> Bool
$c> :: Alias -> Alias -> Bool
> :: Alias -> Alias -> Bool
$c>= :: Alias -> Alias -> Bool
>= :: Alias -> Alias -> Bool
$cmax :: Alias -> Alias -> Alias
max :: Alias -> Alias -> Alias
$cmin :: Alias -> Alias -> Alias
min :: Alias -> Alias -> Alias
Ord, Int -> Alias -> ShowS
[Alias] -> ShowS
Alias -> String
(Int -> Alias -> ShowS)
-> (Alias -> String) -> ([Alias] -> ShowS) -> Show Alias
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Alias -> ShowS
showsPrec :: Int -> Alias -> ShowS
$cshow :: Alias -> String
show :: Alias -> String
$cshowList :: [Alias] -> ShowS
showList :: [Alias] -> ShowS
Show)

instance Pretty Alias where
  pretty :: forall ann. Alias -> Doc ann
pretty (AliasBound VName
v) = VName -> Doc ann
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v
  pretty (AliasFree VName
v) = Doc ann
"~" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> VName -> Doc ann
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v

instance Pretty (S.Set Alias) where
  pretty :: forall ann. Aliases -> Doc ann
pretty = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
braces (Doc ann -> Doc ann) -> (Aliases -> Doc ann) -> Aliases -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
commasep ([Doc ann] -> Doc ann)
-> (Aliases -> [Doc ann]) -> Aliases -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> Doc ann) -> [Alias] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Alias -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Alias -> Doc ann
pretty ([Alias] -> [Doc ann])
-> (Aliases -> [Alias]) -> Aliases -> [Doc ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliases -> [Alias]
forall a. Set a -> [a]
S.toList

-- | The set of in-scope variables that are being aliased.
boundAliases :: Aliases -> S.Set VName
boundAliases :: Aliases -> Names
boundAliases = (Alias -> VName) -> Aliases -> Names
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar (Aliases -> Names) -> (Aliases -> Aliases) -> Aliases -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> Bool) -> Aliases -> Aliases
forall a. (a -> Bool) -> Set a -> Set a
S.filter Alias -> Bool
bound
  where
    bound :: Alias -> Bool
bound AliasBound {} = Bool
True
    bound AliasFree {} = Bool
False

-- | Aliases for a type, which is a set of the variables that are
-- aliased.
type Aliases = S.Set Alias

type TypeAliases = TypeBase Size Aliases

-- | @t \`setAliases\` als@ returns @t@, but with @als@ substituted for
-- any already present aliases.
setAliases :: TypeBase dim asf -> ast -> TypeBase dim ast
setAliases :: forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases TypeBase dim asf
t = TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
addAliases TypeBase dim asf
t ((asf -> ast) -> TypeBase dim ast)
-> (ast -> asf -> ast) -> ast -> TypeBase dim ast
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ast -> asf -> ast
forall a b. a -> b -> a
const

-- | @t \`addAliases\` f@ returns @t@, but with any already present
-- aliases replaced by @f@ applied to that aliases.
addAliases ::
  TypeBase dim asf ->
  (asf -> ast) ->
  TypeBase dim ast
addAliases :: forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
addAliases = ((asf -> ast) -> TypeBase dim asf -> TypeBase dim ast)
-> TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
forall a b c. (a -> b -> c) -> b -> a -> c
flip (asf -> ast) -> TypeBase dim asf -> TypeBase dim ast
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second

aliases :: TypeAliases -> Aliases
aliases :: TypeAliases -> Aliases
aliases = (Size -> Aliases) -> (Aliases -> Aliases) -> TypeAliases -> Aliases
forall m a b. Monoid m => (a -> m) -> (b -> m) -> TypeBase a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap (Aliases -> Size -> Aliases
forall a b. a -> b -> a
const Aliases
forall a. Monoid a => a
mempty) Aliases -> Aliases
forall a. a -> a
id

setFieldAliases :: TypeAliases -> [Name] -> TypeAliases -> TypeAliases
setFieldAliases :: TypeAliases -> [Name] -> TypeAliases -> TypeAliases
setFieldAliases TypeAliases
ve_als (Name
x : [Name]
xs) (Scalar (Record Map Name TypeAliases
fs)) =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Map Name TypeAliases -> ScalarTypeBase Size Aliases
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name TypeAliases -> ScalarTypeBase Size Aliases)
-> Map Name TypeAliases -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ (TypeAliases -> TypeAliases)
-> Name -> Map Name TypeAliases -> Map Name TypeAliases
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (TypeAliases -> [Name] -> TypeAliases -> TypeAliases
setFieldAliases TypeAliases
ve_als [Name]
xs) Name
x Map Name TypeAliases
fs
setFieldAliases TypeAliases
ve_als [Name]
_ TypeAliases
_ = TypeAliases
ve_als

data Entry a
  = Consumable {forall a. Entry a -> a
entryAliases :: a}
  | Nonconsumable {entryAliases :: a}
  deriving (Entry a -> Entry a -> Bool
(Entry a -> Entry a -> Bool)
-> (Entry a -> Entry a -> Bool) -> Eq (Entry a)
forall a. Eq a => Entry a -> Entry a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Entry a -> Entry a -> Bool
== :: Entry a -> Entry a -> Bool
$c/= :: forall a. Eq a => Entry a -> Entry a -> Bool
/= :: Entry a -> Entry a -> Bool
Eq, Eq (Entry a)
Eq (Entry a) =>
(Entry a -> Entry a -> Ordering)
-> (Entry a -> Entry a -> Bool)
-> (Entry a -> Entry a -> Bool)
-> (Entry a -> Entry a -> Bool)
-> (Entry a -> Entry a -> Bool)
-> (Entry a -> Entry a -> Entry a)
-> (Entry a -> Entry a -> Entry a)
-> Ord (Entry a)
Entry a -> Entry a -> Bool
Entry a -> Entry a -> Ordering
Entry a -> Entry a -> Entry a
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
forall a. Ord a => Eq (Entry a)
forall a. Ord a => Entry a -> Entry a -> Bool
forall a. Ord a => Entry a -> Entry a -> Ordering
forall a. Ord a => Entry a -> Entry a -> Entry a
$ccompare :: forall a. Ord a => Entry a -> Entry a -> Ordering
compare :: Entry a -> Entry a -> Ordering
$c< :: forall a. Ord a => Entry a -> Entry a -> Bool
< :: Entry a -> Entry a -> Bool
$c<= :: forall a. Ord a => Entry a -> Entry a -> Bool
<= :: Entry a -> Entry a -> Bool
$c> :: forall a. Ord a => Entry a -> Entry a -> Bool
> :: Entry a -> Entry a -> Bool
$c>= :: forall a. Ord a => Entry a -> Entry a -> Bool
>= :: Entry a -> Entry a -> Bool
$cmax :: forall a. Ord a => Entry a -> Entry a -> Entry a
max :: Entry a -> Entry a -> Entry a
$cmin :: forall a. Ord a => Entry a -> Entry a -> Entry a
min :: Entry a -> Entry a -> Entry a
Ord, Int -> Entry a -> ShowS
[Entry a] -> ShowS
Entry a -> String
(Int -> Entry a -> ShowS)
-> (Entry a -> String) -> ([Entry a] -> ShowS) -> Show (Entry a)
forall a. Show a => Int -> Entry a -> ShowS
forall a. Show a => [Entry a] -> ShowS
forall a. Show a => Entry a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Entry a -> ShowS
showsPrec :: Int -> Entry a -> ShowS
$cshow :: forall a. Show a => Entry a -> String
show :: Entry a -> String
$cshowList :: forall a. Show a => [Entry a] -> ShowS
showList :: [Entry a] -> ShowS
Show)

instance Functor Entry where
  fmap :: forall a b. (a -> b) -> Entry a -> Entry b
fmap a -> b
f (Consumable a
als) = b -> Entry b
forall a. a -> Entry a
Consumable (b -> Entry b) -> b -> Entry b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
als
  fmap a -> b
f (Nonconsumable a
als) = b -> Entry b
forall a. a -> Entry a
Nonconsumable (b -> Entry b) -> b -> Entry b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
als

data CheckEnv = CheckEnv
  { CheckEnv -> Map VName (Entry TypeAliases)
envVtable :: M.Map VName (Entry TypeAliases),
    -- | Location of the definition we are checking.
    CheckEnv -> Loc
envLoc :: Loc
  }

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

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

type Consumed = M.Map VName Loc

data CheckState = CheckState
  { CheckState -> Consumed
stateConsumed :: Consumed,
    CheckState -> DList TypeError
stateErrors :: DL.DList TypeError,
    CheckState -> Map VName NameReason
stateNames :: M.Map VName NameReason,
    CheckState -> Int
stateCounter :: Int
  }

newtype CheckM a = CheckM (ReaderT CheckEnv (State CheckState) a)
  deriving
    ( (forall a b. (a -> b) -> CheckM a -> CheckM b)
-> (forall a b. a -> CheckM b -> CheckM a) -> Functor CheckM
forall a b. a -> CheckM b -> CheckM a
forall a b. (a -> b) -> CheckM a -> CheckM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> CheckM a -> CheckM b
fmap :: forall a b. (a -> b) -> CheckM a -> CheckM b
$c<$ :: forall a b. a -> CheckM b -> CheckM a
<$ :: forall a b. a -> CheckM b -> CheckM a
Functor,
      Functor CheckM
Functor CheckM =>
(forall a. a -> CheckM a)
-> (forall a b. CheckM (a -> b) -> CheckM a -> CheckM b)
-> (forall a b c.
    (a -> b -> c) -> CheckM a -> CheckM b -> CheckM c)
-> (forall a b. CheckM a -> CheckM b -> CheckM b)
-> (forall a b. CheckM a -> CheckM b -> CheckM a)
-> Applicative CheckM
forall a. a -> CheckM a
forall a b. CheckM a -> CheckM b -> CheckM a
forall a b. CheckM a -> CheckM b -> CheckM b
forall a b. CheckM (a -> b) -> CheckM a -> CheckM b
forall a b c. (a -> b -> c) -> CheckM a -> CheckM b -> CheckM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> CheckM a
pure :: forall a. a -> CheckM a
$c<*> :: forall a b. CheckM (a -> b) -> CheckM a -> CheckM b
<*> :: forall a b. CheckM (a -> b) -> CheckM a -> CheckM b
$cliftA2 :: forall a b c. (a -> b -> c) -> CheckM a -> CheckM b -> CheckM c
liftA2 :: forall a b c. (a -> b -> c) -> CheckM a -> CheckM b -> CheckM c
$c*> :: forall a b. CheckM a -> CheckM b -> CheckM b
*> :: forall a b. CheckM a -> CheckM b -> CheckM b
$c<* :: forall a b. CheckM a -> CheckM b -> CheckM a
<* :: forall a b. CheckM a -> CheckM b -> CheckM a
Applicative,
      Applicative CheckM
Applicative CheckM =>
(forall a b. CheckM a -> (a -> CheckM b) -> CheckM b)
-> (forall a b. CheckM a -> CheckM b -> CheckM b)
-> (forall a. a -> CheckM a)
-> Monad CheckM
forall a. a -> CheckM a
forall a b. CheckM a -> CheckM b -> CheckM b
forall a b. CheckM a -> (a -> CheckM b) -> CheckM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. CheckM a -> (a -> CheckM b) -> CheckM b
>>= :: forall a b. CheckM a -> (a -> CheckM b) -> CheckM b
$c>> :: forall a b. CheckM a -> CheckM b -> CheckM b
>> :: forall a b. CheckM a -> CheckM b -> CheckM b
$creturn :: forall a. a -> CheckM a
return :: forall a. a -> CheckM a
Monad,
      MonadReader CheckEnv,
      MonadState CheckState
    )

runCheckM :: Loc -> CheckM a -> (a, [TypeError])
runCheckM :: forall a. Loc -> CheckM a -> (a, [TypeError])
runCheckM Loc
loc (CheckM ReaderT CheckEnv (State CheckState) a
m) =
  let (a
a, CheckState
s) = State CheckState a -> CheckState -> (a, CheckState)
forall s a. State s a -> s -> (a, s)
runState (ReaderT CheckEnv (State CheckState) a
-> CheckEnv -> State CheckState a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT CheckEnv (State CheckState) a
m CheckEnv
env) CheckState
initial_state
   in (a
a, DList TypeError -> [TypeError]
forall a. DList a -> [a]
DL.toList (CheckState -> DList TypeError
stateErrors CheckState
s))
  where
    env :: CheckEnv
env =
      CheckEnv
        { envVtable :: Map VName (Entry TypeAliases)
envVtable = Map VName (Entry TypeAliases)
forall a. Monoid a => a
mempty,
          envLoc :: Loc
envLoc = Loc
loc
        }
    initial_state :: CheckState
initial_state =
      CheckState
        { stateConsumed :: Consumed
stateConsumed = Consumed
forall a. Monoid a => a
mempty,
          stateErrors :: DList TypeError
stateErrors = DList TypeError
forall a. Monoid a => a
mempty,
          stateNames :: Map VName NameReason
stateNames = Map VName NameReason
forall a. Monoid a => a
mempty,
          stateCounter :: Int
stateCounter = Int
0
        }

describeVar :: VName -> CheckM (Doc a)
describeVar :: forall a. VName -> CheckM (Doc a)
describeVar VName
v = do
  loc <- (CheckEnv -> Loc) -> CheckM Loc
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CheckEnv -> Loc
envLoc
  gets $
    maybe ("variable" <+> dquotes (prettyName v)) (nameReason (srclocOf loc))
      . M.lookup v
      . stateNames

noConsumable :: CheckM a -> CheckM a
noConsumable :: forall a. CheckM a -> CheckM a
noConsumable = (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall a. (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((CheckEnv -> CheckEnv) -> CheckM a -> CheckM a)
-> (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall a b. (a -> b) -> a -> b
$ \CheckEnv
env -> CheckEnv
env {envVtable = M.map f $ envVtable env}
  where
    f :: Entry b -> Entry b
f = b -> Entry b
forall a. a -> Entry a
Nonconsumable (b -> Entry b) -> (Entry b -> b) -> Entry b -> Entry b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry b -> b
forall a. Entry a -> a
entryAliases

addError :: (Located loc) => loc -> Notes -> Doc () -> CheckM ()
addError :: forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError loc
loc Notes
notes Doc ()
e = (CheckState -> CheckState) -> CheckM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CheckState -> CheckState) -> CheckM ())
-> (CheckState -> CheckState) -> CheckM ()
forall a b. (a -> b) -> a -> b
$ \CheckState
s ->
  CheckState
s {stateErrors = DL.snoc (stateErrors s) (TypeError (locOf loc) notes e)}

incCounter :: CheckM Int
incCounter :: CheckM Int
incCounter =
  (CheckState -> (Int, CheckState)) -> CheckM Int
forall a. (CheckState -> (a, CheckState)) -> CheckM a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((CheckState -> (Int, CheckState)) -> CheckM Int)
-> (CheckState -> (Int, CheckState)) -> CheckM Int
forall a b. (a -> b) -> a -> b
$ \CheckState
s -> (CheckState -> Int
stateCounter CheckState
s, CheckState
s {stateCounter = stateCounter s + 1})

returnAliased :: Name -> SrcLoc -> CheckM ()
returnAliased :: Name -> SrcLoc -> CheckM ()
returnAliased Name
name SrcLoc
loc =
  SrcLoc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> CheckM ()) -> (Doc () -> Doc ()) -> Doc () -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
withIndexLink Doc ()
"return-aliased" (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
    Doc ()
"Unique-typed return value is aliased to"
      Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (Name -> Doc ()
forall a. Name -> Doc a
forall v a. IsName v => v -> Doc a
prettyName Name
name)
      Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
", which is not consumable."

uniqueReturnAliased :: SrcLoc -> CheckM ()
uniqueReturnAliased :: SrcLoc -> CheckM ()
uniqueReturnAliased SrcLoc
loc =
  SrcLoc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> CheckM ()) -> (Doc () -> Doc ()) -> Doc () -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
withIndexLink Doc ()
"unique-return-aliased" (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
    Doc ()
"A unique-typed component of the return value is aliased to some other component."

checkReturnAlias :: SrcLoc -> [Pat ParamType] -> ResType -> TypeAliases -> CheckM ()
checkReturnAlias :: SrcLoc
-> [Pat (TypeBase Size Diet)]
-> ResType
-> TypeAliases
-> CheckM ()
checkReturnAlias SrcLoc
loc [Pat (TypeBase Size Diet)]
params ResType
rettp =
  (Set (Uniqueness, VName)
 -> (Uniqueness, Names) -> CheckM (Set (Uniqueness, VName)))
-> Set (Uniqueness, VName) -> [(Uniqueness, Names)] -> CheckM ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ ([Pat (TypeBase Size Diet)]
-> Set (Uniqueness, VName)
-> (Uniqueness, Names)
-> CheckM (Set (Uniqueness, VName))
forall {t :: * -> *} {dim}.
Foldable t =>
t (Pat (TypeBase dim Diet))
-> Set (Uniqueness, VName)
-> (Uniqueness, Names)
-> CheckM (Set (Uniqueness, VName))
checkReturnAlias' [Pat (TypeBase Size Diet)]
params) Set (Uniqueness, VName)
forall a. Set a
S.empty ([(Uniqueness, Names)] -> CheckM ())
-> (TypeAliases -> [(Uniqueness, Names)])
-> TypeAliases
-> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ResType -> TypeAliases -> [(Uniqueness, Names)]
forall {shape}.
TypeBase shape Uniqueness -> TypeAliases -> [(Uniqueness, Names)]
returnAliases ResType
rettp
  where
    checkReturnAlias' :: t (Pat (TypeBase dim Diet))
-> Set (Uniqueness, VName)
-> (Uniqueness, Names)
-> CheckM (Set (Uniqueness, VName))
checkReturnAlias' t (Pat (TypeBase dim Diet))
params' Set (Uniqueness, VName)
seen (Uniqueness
Unique, Names
names) = do
      Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` ((Uniqueness, VName) -> VName) -> Set (Uniqueness, VName) -> Names
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map (Uniqueness, VName) -> VName
forall a b. (a, b) -> b
snd Set (Uniqueness, VName)
seen) ([VName] -> Bool) -> [VName] -> Bool
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
forall a. Set a -> [a]
S.toList Names
names) (CheckM () -> CheckM ()) -> CheckM () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
        SrcLoc -> CheckM ()
uniqueReturnAliased SrcLoc
loc
      t (Pat (TypeBase dim Diet)) -> Names -> CheckM ()
forall {t :: * -> *} {dim}.
Foldable t =>
t (Pat (TypeBase dim Diet)) -> Names -> CheckM ()
notAliasesParam t (Pat (TypeBase dim Diet))
params' Names
names
      Set (Uniqueness, VName) -> CheckM (Set (Uniqueness, VName))
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set (Uniqueness, VName) -> CheckM (Set (Uniqueness, VName)))
-> Set (Uniqueness, VName) -> CheckM (Set (Uniqueness, VName))
forall a b. (a -> b) -> a -> b
$ Set (Uniqueness, VName)
seen Set (Uniqueness, VName)
-> Set (Uniqueness, VName) -> Set (Uniqueness, VName)
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Uniqueness -> Names -> Set (Uniqueness, VName)
forall {t} {a}. (Ord t, Ord a) => t -> Set a -> Set (t, a)
tag Uniqueness
Unique Names
names
    checkReturnAlias' t (Pat (TypeBase dim Diet))
_ Set (Uniqueness, VName)
seen (Uniqueness
Nonunique, Names
names) = do
      Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (((Uniqueness, VName) -> Bool) -> [(Uniqueness, VName)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Uniqueness, VName) -> Set (Uniqueness, VName) -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set (Uniqueness, VName)
seen) ([(Uniqueness, VName)] -> Bool) -> [(Uniqueness, VName)] -> Bool
forall a b. (a -> b) -> a -> b
$ Set (Uniqueness, VName) -> [(Uniqueness, VName)]
forall a. Set a -> [a]
S.toList (Set (Uniqueness, VName) -> [(Uniqueness, VName)])
-> Set (Uniqueness, VName) -> [(Uniqueness, VName)]
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Names -> Set (Uniqueness, VName)
forall {t} {a}. (Ord t, Ord a) => t -> Set a -> Set (t, a)
tag Uniqueness
Unique Names
names) (CheckM () -> CheckM ()) -> CheckM () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
        SrcLoc -> CheckM ()
uniqueReturnAliased SrcLoc
loc
      Set (Uniqueness, VName) -> CheckM (Set (Uniqueness, VName))
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set (Uniqueness, VName) -> CheckM (Set (Uniqueness, VName)))
-> Set (Uniqueness, VName) -> CheckM (Set (Uniqueness, VName))
forall a b. (a -> b) -> a -> b
$ Set (Uniqueness, VName)
seen Set (Uniqueness, VName)
-> Set (Uniqueness, VName) -> Set (Uniqueness, VName)
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Uniqueness -> Names -> Set (Uniqueness, VName)
forall {t} {a}. (Ord t, Ord a) => t -> Set a -> Set (t, a)
tag Uniqueness
Nonunique Names
names

    notAliasesParam :: t (Pat (TypeBase dim Diet)) -> Names -> CheckM ()
notAliasesParam t (Pat (TypeBase dim Diet))
params' Names
names =
      t (Pat (TypeBase dim Diet))
-> (Pat (TypeBase dim Diet) -> CheckM ()) -> CheckM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t (Pat (TypeBase dim Diet))
params' ((Pat (TypeBase dim Diet) -> CheckM ()) -> CheckM ())
-> (Pat (TypeBase dim Diet) -> CheckM ()) -> CheckM ()
forall a b. (a -> b) -> a -> b
$ \Pat (TypeBase dim Diet)
p ->
        let consumedNonunique :: (VName, TypeBase dim Diet) -> Bool
consumedNonunique (VName
v, TypeBase dim Diet
t) =
              Bool -> Bool
not (TypeBase dim Diet -> Bool
forall {dim}. TypeBase dim Diet -> Bool
consumableParamType TypeBase dim Diet
t) Bool -> Bool -> Bool
&& (VName
v VName -> Names -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Names
names)
         in case ((VName, TypeBase dim Diet) -> Bool)
-> [(VName, TypeBase dim Diet)] -> Maybe (VName, TypeBase dim Diet)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (VName, TypeBase dim Diet) -> Bool
forall {dim}. (VName, TypeBase dim Diet) -> Bool
consumedNonunique ([(VName, TypeBase dim Diet)] -> Maybe (VName, TypeBase dim Diet))
-> [(VName, TypeBase dim Diet)] -> Maybe (VName, TypeBase dim Diet)
forall a b. (a -> b) -> a -> b
$ Pat (TypeBase dim Diet) -> [(VName, TypeBase dim Diet)]
forall t. Pat t -> [(VName, t)]
patternMap Pat (TypeBase dim Diet)
p of
              Just (VName
v, TypeBase dim Diet
_) ->
                Name -> SrcLoc -> CheckM ()
returnAliased (VName -> Name
baseName VName
v) SrcLoc
loc
              Maybe (VName, TypeBase dim Diet)
Nothing ->
                () -> CheckM ()
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    tag :: t -> Set a -> Set (t, a)
tag t
u = (a -> (t, a)) -> Set a -> Set (t, a)
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map (t
u,)

    returnAliases :: TypeBase shape Uniqueness -> TypeAliases -> [(Uniqueness, Names)]
returnAliases (Scalar (Record Map Name (TypeBase shape Uniqueness)
ets1)) (Scalar (Record Map Name TypeAliases
ets2)) =
      [[(Uniqueness, Names)]] -> [(Uniqueness, Names)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Uniqueness, Names)]] -> [(Uniqueness, Names)])
-> [[(Uniqueness, Names)]] -> [(Uniqueness, Names)]
forall a b. (a -> b) -> a -> b
$ Map Name [(Uniqueness, Names)] -> [[(Uniqueness, Names)]]
forall k a. Map k a -> [a]
M.elems (Map Name [(Uniqueness, Names)] -> [[(Uniqueness, Names)]])
-> Map Name [(Uniqueness, Names)] -> [[(Uniqueness, Names)]]
forall a b. (a -> b) -> a -> b
$ (TypeBase shape Uniqueness -> TypeAliases -> [(Uniqueness, Names)])
-> Map Name (TypeBase shape Uniqueness)
-> Map Name TypeAliases
-> Map Name [(Uniqueness, Names)]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase shape Uniqueness -> TypeAliases -> [(Uniqueness, Names)]
returnAliases Map Name (TypeBase shape Uniqueness)
ets1 Map Name TypeAliases
ets2
    returnAliases TypeBase shape Uniqueness
expected TypeAliases
got =
      [(TypeBase shape Uniqueness -> Uniqueness
forall shape. TypeBase shape Uniqueness -> Uniqueness
uniqueness TypeBase shape Uniqueness
expected, (Alias -> VName) -> Aliases -> Names
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar (Aliases -> Names) -> Aliases -> Names
forall a b. (a -> b) -> a -> b
$ TypeAliases -> Aliases
aliases TypeAliases
got)]

    consumableParamType :: TypeBase dim Diet -> Bool
consumableParamType (Array Diet
u Shape dim
_ ScalarTypeBase dim NoUniqueness
_) = Diet
u Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume
    consumableParamType (Scalar Prim {}) = Bool
True
    consumableParamType (Scalar (TypeVar Diet
u QualName VName
_ [TypeArg dim]
_)) = Diet
u Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume
    consumableParamType (Scalar (Record Map Name (TypeBase dim Diet)
fs)) = (TypeBase dim Diet -> Bool) -> Map Name (TypeBase dim Diet) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim Diet -> Bool
consumableParamType Map Name (TypeBase dim Diet)
fs
    consumableParamType (Scalar (Sum Map Name [TypeBase dim Diet]
fs)) = ([TypeBase dim Diet] -> Bool)
-> Map Name [TypeBase dim Diet] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((TypeBase dim Diet -> Bool) -> [TypeBase dim Diet] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim Diet -> Bool
consumableParamType) Map Name [TypeBase dim Diet]
fs
    consumableParamType (Scalar Arrow {}) = Bool
False

unscope :: [VName] -> Aliases -> Aliases
unscope :: [VName] -> Aliases -> Aliases
unscope [VName]
bound = (Alias -> Alias) -> Aliases -> Aliases
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> Alias
f
  where
    f :: Alias -> Alias
f (AliasFree VName
v) = VName -> Alias
AliasFree VName
v
    f (AliasBound VName
v) = if VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
bound then VName -> Alias
AliasFree VName
v else VName -> Alias
AliasBound VName
v

-- | Figure out the aliases of each bound name in a pattern.
matchPat :: Pat t -> TypeAliases -> DL.DList (VName, (t, TypeAliases))
matchPat :: forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat (PatParens PatBase Info VName t
p SrcLoc
_) TypeAliases
t = PatBase Info VName t
-> TypeAliases -> DList (VName, (t, TypeAliases))
forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat PatBase Info VName t
p TypeAliases
t
matchPat (TuplePat [PatBase Info VName t]
ps SrcLoc
_) TypeAliases
t
  | Just [TypeAliases]
ts <- TypeAliases -> Maybe [TypeAliases]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord TypeAliases
t = [DList (VName, (t, TypeAliases))]
-> DList (VName, (t, TypeAliases))
forall a. Monoid a => [a] -> a
mconcat ([DList (VName, (t, TypeAliases))]
 -> DList (VName, (t, TypeAliases)))
-> [DList (VName, (t, TypeAliases))]
-> DList (VName, (t, TypeAliases))
forall a b. (a -> b) -> a -> b
$ (PatBase Info VName t
 -> TypeAliases -> DList (VName, (t, TypeAliases)))
-> [PatBase Info VName t]
-> [TypeAliases]
-> [DList (VName, (t, TypeAliases))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName t
-> TypeAliases -> DList (VName, (t, TypeAliases))
forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat [PatBase Info VName t]
ps [TypeAliases]
ts
matchPat (RecordPat [(L Name, PatBase Info VName t)]
fs1 SrcLoc
_) (Scalar (Record Map Name TypeAliases
fs2)) =
  [DList (VName, (t, TypeAliases))]
-> DList (VName, (t, TypeAliases))
forall a. Monoid a => [a] -> a
mconcat ([DList (VName, (t, TypeAliases))]
 -> DList (VName, (t, TypeAliases)))
-> [DList (VName, (t, TypeAliases))]
-> DList (VName, (t, TypeAliases))
forall a b. (a -> b) -> a -> b
$
    (PatBase Info VName t
 -> TypeAliases -> DList (VName, (t, TypeAliases)))
-> [PatBase Info VName t]
-> [TypeAliases]
-> [DList (VName, (t, TypeAliases))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
      PatBase Info VName t
-> TypeAliases -> DList (VName, (t, TypeAliases))
forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat
      (((Name, PatBase Info VName t) -> PatBase Info VName t)
-> [(Name, PatBase Info VName t)] -> [PatBase Info VName t]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatBase Info VName t) -> PatBase Info VName t
forall a b. (a, b) -> b
snd (Map Name (PatBase Info VName t) -> [(Name, PatBase Info VName t)]
forall a. Map Name a -> [(Name, a)]
sortFields ([(Name, PatBase Info VName t)] -> Map Name (PatBase Info VName t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (((L Name, PatBase Info VName t) -> (Name, PatBase Info VName t))
-> [(L Name, PatBase Info VName t)]
-> [(Name, PatBase Info VName t)]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, PatBase Info VName t) -> (Name, PatBase Info VName t)
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 L Name -> Name
forall a. L a -> a
unLoc) [(L Name, PatBase Info VName t)]
fs1))))
      (((Name, TypeAliases) -> TypeAliases)
-> [(Name, TypeAliases)] -> [TypeAliases]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeAliases) -> TypeAliases
forall a b. (a, b) -> b
snd (Map Name TypeAliases -> [(Name, TypeAliases)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name TypeAliases
fs2))
matchPat (Id VName
v (Info t
t) SrcLoc
_) TypeAliases
als = (VName, (t, TypeAliases)) -> DList (VName, (t, TypeAliases))
forall a. a -> DList a
DL.singleton (VName
v, (t
t, TypeAliases
als))
matchPat (PatAscription PatBase Info VName t
p TypeExp Size VName
_ SrcLoc
_) TypeAliases
t = PatBase Info VName t
-> TypeAliases -> DList (VName, (t, TypeAliases))
forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat PatBase Info VName t
p TypeAliases
t
matchPat (PatConstr Name
v Info t
_ [PatBase Info VName t]
ps SrcLoc
_) (Scalar (Sum Map Name [TypeAliases]
cs))
  | Just [TypeAliases]
ts <- Name -> Map Name [TypeAliases] -> Maybe [TypeAliases]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
v Map Name [TypeAliases]
cs = [DList (VName, (t, TypeAliases))]
-> DList (VName, (t, TypeAliases))
forall a. Monoid a => [a] -> a
mconcat ([DList (VName, (t, TypeAliases))]
 -> DList (VName, (t, TypeAliases)))
-> [DList (VName, (t, TypeAliases))]
-> DList (VName, (t, TypeAliases))
forall a b. (a -> b) -> a -> b
$ (PatBase Info VName t
 -> TypeAliases -> DList (VName, (t, TypeAliases)))
-> [PatBase Info VName t]
-> [TypeAliases]
-> [DList (VName, (t, TypeAliases))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName t
-> TypeAliases -> DList (VName, (t, TypeAliases))
forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat [PatBase Info VName t]
ps [TypeAliases]
ts
matchPat TuplePat {} TypeAliases
_ = DList (VName, (t, TypeAliases))
forall a. Monoid a => a
mempty
matchPat RecordPat {} TypeAliases
_ = DList (VName, (t, TypeAliases))
forall a. Monoid a => a
mempty
matchPat PatConstr {} TypeAliases
_ = DList (VName, (t, TypeAliases))
forall a. Monoid a => a
mempty
matchPat Wildcard {} TypeAliases
_ = DList (VName, (t, TypeAliases))
forall a. Monoid a => a
mempty
matchPat PatLit {} TypeAliases
_ = DList (VName, (t, TypeAliases))
forall a. Monoid a => a
mempty
matchPat (PatAttr AttrInfo VName
_ PatBase Info VName t
p SrcLoc
_) TypeAliases
t = PatBase Info VName t
-> TypeAliases -> DList (VName, (t, TypeAliases))
forall t. Pat t -> TypeAliases -> DList (VName, (t, TypeAliases))
matchPat PatBase Info VName t
p TypeAliases
t

bindingPat ::
  Pat StructType ->
  TypeAliases ->
  CheckM (a, TypeAliases) ->
  CheckM (a, TypeAliases)
bindingPat :: forall a.
Pat StructType
-> TypeAliases
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
bindingPat Pat StructType
p TypeAliases
t = ((a, TypeAliases) -> (a, TypeAliases))
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a b. (a -> b) -> CheckM a -> CheckM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeAliases -> TypeAliases)
-> (a, TypeAliases) -> (a, TypeAliases)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Aliases -> Aliases) -> TypeAliases -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([VName] -> Aliases -> Aliases
unscope (Pat StructType -> [VName]
forall t. Pat t -> [VName]
patNames Pat StructType
p)))) (CheckM (a, TypeAliases) -> CheckM (a, TypeAliases))
-> (CheckM (a, TypeAliases) -> CheckM (a, TypeAliases))
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CheckEnv -> CheckEnv)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a. (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local CheckEnv -> CheckEnv
bind
  where
    bind :: CheckEnv -> CheckEnv
bind CheckEnv
env =
      CheckEnv
env
        { envVtable =
            foldr (uncurry M.insert . f) (envVtable env) (matchPat p t)
        }
      where
        f :: (VName, (a, p a Aliases)) -> (VName, Entry (p a Aliases))
f (VName
v, (a
_, p a Aliases
als)) = (VName
v, p a Aliases -> Entry (p a Aliases)
forall a. a -> Entry a
Consumable (p a Aliases -> Entry (p a Aliases))
-> p a Aliases -> Entry (p a Aliases)
forall a b. (a -> b) -> a -> b
$ (Aliases -> Aliases) -> p a Aliases -> p a Aliases
forall b c a. (b -> c) -> p a b -> p a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Alias -> Aliases -> Aliases
forall a. Ord a => a -> Set a -> Set a
S.insert (VName -> Alias
AliasBound VName
v)) p a Aliases
als)

bindingParam :: Pat ParamType -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParam :: forall a.
Pat (TypeBase Size Diet)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParam Pat (TypeBase Size Diet)
p CheckM (a, TypeAliases)
m = do
  (TypeBase Size Diet -> CheckM ())
-> Pat (TypeBase Size Diet) -> CheckM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (CheckM () -> CheckM ()
forall a. CheckM a -> CheckM a
noConsumable (CheckM () -> CheckM ())
-> (TypeBase Size Diet -> CheckM ())
-> TypeBase Size Diet
-> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Size -> CheckM (Size, TypeAliases))
-> (Diet -> CheckM Diet) -> TypeBase Size Diet -> CheckM ()
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bifoldable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f ()
bitraverse_ Size -> CheckM (Size, TypeAliases)
checkExp Diet -> CheckM Diet
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) Pat (TypeBase Size Diet)
p
  (TypeAliases -> TypeAliases)
-> (a, TypeAliases) -> (a, TypeAliases)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Aliases -> Aliases) -> TypeAliases -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([VName] -> Aliases -> Aliases
unscope (Pat (TypeBase Size Diet) -> [VName]
forall t. Pat t -> [VName]
patNames Pat (TypeBase Size Diet)
p))) ((a, TypeAliases) -> (a, TypeAliases))
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CheckEnv -> CheckEnv)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a. (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local CheckEnv -> CheckEnv
bind CheckM (a, TypeAliases)
m
  where
    bind :: CheckEnv -> CheckEnv
bind CheckEnv
env =
      CheckEnv
env
        { envVtable =
            foldr (uncurry M.insert . f) (envVtable env) (patternMap p)
        }
    f :: (VName, TypeBase dim Diet) -> (VName, Entry (TypeBase dim Aliases))
f (VName
v, TypeBase dim Diet
t)
      | TypeBase dim Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase dim Diet
t Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume = (VName
v, TypeBase dim Aliases -> Entry (TypeBase dim Aliases)
forall a. a -> Entry a
Consumable (TypeBase dim Aliases -> Entry (TypeBase dim Aliases))
-> TypeBase dim Aliases -> Entry (TypeBase dim Aliases)
forall a b. (a -> b) -> a -> b
$ TypeBase dim Diet
t TypeBase dim Diet -> Aliases -> TypeBase dim Aliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Alias -> Aliases
forall a. a -> Set a
S.singleton (VName -> Alias
AliasBound VName
v))
      | Bool
otherwise = (VName
v, TypeBase dim Aliases -> Entry (TypeBase dim Aliases)
forall a. a -> Entry a
Nonconsumable (TypeBase dim Aliases -> Entry (TypeBase dim Aliases))
-> TypeBase dim Aliases -> Entry (TypeBase dim Aliases)
forall a b. (a -> b) -> a -> b
$ TypeBase dim Diet
t TypeBase dim Diet -> Aliases -> TypeBase dim Aliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Alias -> Aliases
forall a. a -> Set a
S.singleton (VName -> Alias
AliasBound VName
v))

bindingIdent :: Diet -> Ident StructType -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingIdent :: forall a.
Diet
-> IdentBase Info VName StructType
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
bindingIdent Diet
d (Ident VName
v (Info StructType
t) SrcLoc
_) =
  ((a, TypeAliases) -> (a, TypeAliases))
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a b. (a -> b) -> CheckM a -> CheckM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeAliases -> TypeAliases)
-> (a, TypeAliases) -> (a, TypeAliases)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Aliases -> Aliases) -> TypeAliases -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([VName] -> Aliases -> Aliases
unscope [VName
v]))) (CheckM (a, TypeAliases) -> CheckM (a, TypeAliases))
-> (CheckM (a, TypeAliases) -> CheckM (a, TypeAliases))
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CheckEnv -> CheckEnv)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a. (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local CheckEnv -> CheckEnv
bind
  where
    bind :: CheckEnv -> CheckEnv
bind CheckEnv
env = CheckEnv
env {envVtable = M.insert v t' (envVtable env)}
    d' :: a -> Entry a
d' = case Diet
d of
      Diet
Consume -> a -> Entry a
forall a. a -> Entry a
Consumable
      Diet
Observe -> a -> Entry a
forall a. a -> Entry a
Nonconsumable
    t' :: Entry TypeAliases
t' = TypeAliases -> Entry TypeAliases
forall a. a -> Entry a
d' (TypeAliases -> Entry TypeAliases)
-> TypeAliases -> Entry TypeAliases
forall a b. (a -> b) -> a -> b
$ StructType
t StructType -> Aliases -> TypeAliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Alias -> Aliases
forall a. a -> Set a
S.singleton (VName -> Alias
AliasBound VName
v)

bindingParams :: [Pat ParamType] -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParams :: forall a.
[Pat (TypeBase Size Diet)]
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParams [Pat (TypeBase Size Diet)]
params CheckM (a, TypeAliases)
m =
  CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a. CheckM a -> CheckM a
noConsumable (CheckM (a, TypeAliases) -> CheckM (a, TypeAliases))
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a b. (a -> b) -> a -> b
$
    (TypeAliases -> TypeAliases)
-> (a, TypeAliases) -> (a, TypeAliases)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Aliases -> Aliases) -> TypeAliases -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([VName] -> Aliases -> Aliases
unscope ((Pat (TypeBase Size Diet) -> [VName])
-> [Pat (TypeBase Size Diet)] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat (TypeBase Size Diet) -> [VName]
forall t. Pat t -> [VName]
patNames [Pat (TypeBase Size Diet)]
params)))
      ((a, TypeAliases) -> (a, TypeAliases))
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat (TypeBase Size Diet)
 -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases))
-> CheckM (a, TypeAliases)
-> [Pat (TypeBase Size Diet)]
-> CheckM (a, TypeAliases)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Pat (TypeBase Size Diet)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a.
Pat (TypeBase Size Diet)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParam CheckM (a, TypeAliases)
m [Pat (TypeBase Size Diet)]
params

bindingLoopForm :: LoopFormBase Info VName -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingLoopForm :: forall a.
LoopFormBase Info VName
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingLoopForm (For IdentBase Info VName StructType
ident Size
_) CheckM (a, TypeAliases)
m = Diet
-> IdentBase Info VName StructType
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
forall a.
Diet
-> IdentBase Info VName StructType
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
bindingIdent Diet
Observe IdentBase Info VName StructType
ident CheckM (a, TypeAliases)
m
bindingLoopForm (ForIn Pat StructType
pat Size
_) CheckM (a, TypeAliases)
m = Pat (TypeBase Size Diet)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
forall a.
Pat (TypeBase Size Diet)
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParam Pat (TypeBase Size Diet)
pat' CheckM (a, TypeAliases)
m
  where
    pat' :: Pat (TypeBase Size Diet)
pat' = (StructType -> TypeBase Size Diet)
-> Pat StructType -> Pat (TypeBase Size Diet)
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((NoUniqueness -> Diet) -> StructType -> TypeBase Size Diet
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Diet -> NoUniqueness -> Diet
forall a b. a -> b -> a
const Diet
Observe)) Pat StructType
pat
bindingLoopForm While {} CheckM (a, TypeAliases)
m = CheckM (a, TypeAliases)
m

bindingFun :: VName -> TypeAliases -> CheckM a -> CheckM a
bindingFun :: forall a. VName -> TypeAliases -> CheckM a -> CheckM a
bindingFun VName
v TypeAliases
t = (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall a. (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((CheckEnv -> CheckEnv) -> CheckM a -> CheckM a)
-> (CheckEnv -> CheckEnv) -> CheckM a -> CheckM a
forall a b. (a -> b) -> a -> b
$ \CheckEnv
env ->
  CheckEnv
env {envVtable = M.insert v (Nonconsumable t) (envVtable env)}

checkIfConsumed :: Loc -> Aliases -> CheckM ()
checkIfConsumed :: Loc -> Aliases -> CheckM ()
checkIfConsumed Loc
rloc Aliases
als = do
  cons <- (CheckState -> Consumed) -> CheckM Consumed
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CheckState -> Consumed
stateConsumed
  let bad VName
v = (Loc -> (VName, Loc)) -> Maybe Loc -> Maybe (VName, Loc)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (VName
v,) (Maybe Loc -> Maybe (VName, Loc))
-> Maybe Loc -> Maybe (VName, Loc)
forall a b. (a -> b) -> a -> b
$ VName
v VName -> Consumed -> Maybe Loc
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Consumed
cons
  forM_ (mapMaybe (bad . aliasVar) $ S.toList als) $ \(VName
v, Loc
wloc) -> do
    v' <- VName -> CheckM (Doc ())
forall a. VName -> CheckM (Doc a)
describeVar VName
v
    addError rloc mempty . withIndexLink "use-after-consume" $
      "Using"
        <+> v'
        <> ", but this was consumed at"
          <+> pretty (locStrRel rloc wloc)
        <> ".  (Possibly through aliases.)"

consumed :: Consumed -> CheckM ()
consumed :: Consumed -> CheckM ()
consumed Consumed
vs = (CheckState -> CheckState) -> CheckM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CheckState -> CheckState) -> CheckM ())
-> (CheckState -> CheckState) -> CheckM ()
forall a b. (a -> b) -> a -> b
$ \CheckState
s -> CheckState
s {stateConsumed = stateConsumed s <> vs}

consumeAliases :: Loc -> Aliases -> CheckM ()
consumeAliases :: Loc -> Aliases -> CheckM ()
consumeAliases Loc
loc Aliases
als = do
  vtable <- (CheckEnv -> Map VName (Entry TypeAliases))
-> CheckM (Map VName (Entry TypeAliases))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CheckEnv -> Map VName (Entry TypeAliases)
envVtable
  let isBad VName
v =
        case VName
v VName -> Map VName (Entry TypeAliases) -> Maybe (Entry TypeAliases)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Entry TypeAliases)
vtable of
          Just (Nonconsumable {}) -> Bool
True
          Just Entry TypeAliases
_ -> Bool
False
          Maybe (Entry TypeAliases)
Nothing -> Bool
True
      checkIfConsumable (AliasBound VName
v)
        | VName -> Bool
isBad VName
v = do
            v' <- VName -> CheckM (Doc ())
forall a. VName -> CheckM (Doc a)
describeVar VName
v
            addError loc mempty . withIndexLink "not-consumable" $
              "Consuming" <+> v' <> ", which is not consumable."
      checkIfConsumable Alias
_ = () -> CheckM ()
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  mapM_ checkIfConsumable $ S.toList als
  checkIfConsumed loc als
  consumed als'
  where
    als' :: Consumed
als' = [(VName, Loc)] -> Consumed
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Loc)] -> Consumed) -> [(VName, Loc)] -> Consumed
forall a b. (a -> b) -> a -> b
$ (Alias -> (VName, Loc)) -> [Alias] -> [(VName, Loc)]
forall a b. (a -> b) -> [a] -> [b]
map ((,Loc
loc) (VName -> (VName, Loc))
-> (Alias -> VName) -> Alias -> (VName, Loc)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) ([Alias] -> [(VName, Loc)]) -> [Alias] -> [(VName, Loc)]
forall a b. (a -> b) -> a -> b
$ Aliases -> [Alias]
forall a. Set a -> [a]
S.toList Aliases
als

consume :: Loc -> VName -> StructType -> CheckM ()
consume :: Loc -> VName -> StructType -> CheckM ()
consume Loc
loc VName
v StructType
t =
  Loc -> Aliases -> CheckM ()
consumeAliases Loc
loc (Aliases -> CheckM ())
-> (TypeAliases -> Aliases) -> TypeAliases -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeAliases -> Aliases
aliases (TypeAliases -> CheckM ()) -> CheckM TypeAliases -> CheckM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Loc -> VName -> StructType -> CheckM TypeAliases
observeVar Loc
loc VName
v StructType
t

-- | Observe the given name here and return its aliases.
observeVar :: Loc -> VName -> StructType -> CheckM TypeAliases
observeVar :: Loc -> VName -> StructType -> CheckM TypeAliases
observeVar Loc
loc VName
v StructType
t = do
  als <-
    (CheckEnv -> TypeAliases) -> CheckM TypeAliases
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((CheckEnv -> TypeAliases) -> CheckM TypeAliases)
-> (CheckEnv -> TypeAliases) -> CheckM TypeAliases
forall a b. (a -> b) -> a -> b
$ \CheckEnv
env ->
      TypeAliases
-> (Entry TypeAliases -> TypeAliases)
-> Maybe (Entry TypeAliases)
-> TypeAliases
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map VName (Entry TypeAliases) -> TypeAliases
forall {a}. Map VName a -> TypeAliases
isGlobal (CheckEnv -> Map VName (Entry TypeAliases)
envVtable CheckEnv
env)) Entry TypeAliases -> TypeAliases
forall a. Entry a -> a
isLocal (Maybe (Entry TypeAliases) -> TypeAliases)
-> Maybe (Entry TypeAliases) -> TypeAliases
forall a b. (a -> b) -> a -> b
$
        VName -> Map VName (Entry TypeAliases) -> Maybe (Entry TypeAliases)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (CheckEnv -> Map VName (Entry TypeAliases)
envVtable CheckEnv
env)
  checkIfConsumed loc (aliases als)
  pure als
  where
    isLocal :: Entry a -> a
isLocal = Entry a -> a
forall a. Entry a -> a
entryAliases

    -- Handling globals is tricky.  For arrays and such, we do want to
    -- track their aliases.  We do not want to track the aliases of
    -- functions.  However, array bindings that are *polymorphic*
    -- should be treated like functions.  However, we do not have
    -- access to the original binding information here.  To avoid
    -- having to plumb that all the way here, we infer that an array
    -- binding is a polymorphic instantiation if its size contains any
    -- locally bound names.
    isGlobal :: Map VName a -> TypeAliases
isGlobal Map VName a
vtable
      | Map VName a -> StructType -> Bool
forall {a} {u}. Map VName a -> TypeBase Size u -> Bool
isInstantiation Map VName a
vtable StructType
t = (NoUniqueness -> Aliases) -> StructType -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliases -> NoUniqueness -> Aliases
forall a b. a -> b -> a
const Aliases
forall a. Monoid a => a
mempty) StructType
t
      | Bool
otherwise = TypeAliases -> TypeAliases
forall {dim}. TypeBase dim Aliases -> TypeBase dim Aliases
selfAlias (TypeAliases -> TypeAliases) -> TypeAliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ (NoUniqueness -> Aliases) -> StructType -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliases -> NoUniqueness -> Aliases
forall a b. a -> b -> a
const Aliases
forall a. Monoid a => a
mempty) StructType
t

    isInstantiation :: Map VName a -> TypeBase Size u -> Bool
isInstantiation Map VName a
vtable =
      (VName -> Bool) -> Names -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Map VName a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName a
vtable) (Names -> Bool)
-> (TypeBase Size u -> Names) -> TypeBase Size u -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Names
fvVars (FV -> Names)
-> (TypeBase Size u -> FV) -> TypeBase Size u -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Size u -> FV
forall u. TypeBase Size u -> FV
freeInType

    selfAlias :: TypeBase dim Aliases -> TypeBase dim Aliases
selfAlias (Array Aliases
als Shape dim
shape ScalarTypeBase dim NoUniqueness
et) = Aliases
-> Shape dim
-> ScalarTypeBase dim NoUniqueness
-> TypeBase dim Aliases
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array (Alias -> Aliases -> Aliases
forall a. Ord a => a -> Set a -> Set a
S.insert (VName -> Alias
AliasBound VName
v) Aliases
als) Shape dim
shape ScalarTypeBase dim NoUniqueness
et
    selfAlias (Scalar ScalarTypeBase dim Aliases
st) = ScalarTypeBase dim Aliases -> TypeBase dim Aliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim Aliases -> TypeBase dim Aliases)
-> ScalarTypeBase dim Aliases -> TypeBase dim Aliases
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim Aliases -> ScalarTypeBase dim Aliases
selfAlias' ScalarTypeBase dim Aliases
st
    selfAlias' :: ScalarTypeBase dim Aliases -> ScalarTypeBase dim Aliases
selfAlias' (TypeVar Aliases
als QualName VName
tn [TypeArg dim]
args) = Aliases
-> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim Aliases
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar Aliases
als QualName VName
tn [TypeArg dim]
args -- #1675 FIXME
    selfAlias' (Record Map Name (TypeBase dim Aliases)
fs) = Map Name (TypeBase dim Aliases) -> ScalarTypeBase dim Aliases
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase dim Aliases) -> ScalarTypeBase dim Aliases)
-> Map Name (TypeBase dim Aliases) -> ScalarTypeBase dim Aliases
forall a b. (a -> b) -> a -> b
$ (TypeBase dim Aliases -> TypeBase dim Aliases)
-> Map Name (TypeBase dim Aliases)
-> Map Name (TypeBase dim Aliases)
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim Aliases -> TypeBase dim Aliases
selfAlias Map Name (TypeBase dim Aliases)
fs
    selfAlias' (Sum Map Name [TypeBase dim Aliases]
fs) = Map Name [TypeBase dim Aliases] -> ScalarTypeBase dim Aliases
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase dim Aliases] -> ScalarTypeBase dim Aliases)
-> Map Name [TypeBase dim Aliases] -> ScalarTypeBase dim Aliases
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim Aliases] -> [TypeBase dim Aliases])
-> Map Name [TypeBase dim Aliases]
-> Map Name [TypeBase dim Aliases]
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim Aliases -> TypeBase dim Aliases)
-> [TypeBase dim Aliases] -> [TypeBase dim Aliases]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase dim Aliases -> TypeBase dim Aliases
selfAlias) Map Name [TypeBase dim Aliases]
fs
    selfAlias' et :: ScalarTypeBase dim Aliases
et@Arrow {} = ScalarTypeBase dim Aliases
et
    selfAlias' et :: ScalarTypeBase dim Aliases
et@Prim {} = ScalarTypeBase dim Aliases
et

-- Capture any newly consumed variables that occur during the provided action.
contain :: CheckM a -> CheckM (a, Consumed)
contain :: forall a. CheckM a -> CheckM (a, Consumed)
contain CheckM a
m = do
  prev_cons <- (CheckState -> Consumed) -> CheckM Consumed
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CheckState -> Consumed
stateConsumed
  x <- m
  new_cons <- gets $ (`M.difference` prev_cons) . stateConsumed
  modify $ \CheckState
s -> CheckState
s {stateConsumed = prev_cons}
  pure (x, new_cons)

-- | The two types are assumed to be approximately structurally equal,
-- but not necessarily regarding sizes.  Combines aliases and prefers
-- other information from first argument.
combineAliases :: TypeAliases -> TypeAliases -> TypeAliases
combineAliases :: TypeAliases -> TypeAliases -> TypeAliases
combineAliases (Array Aliases
als1 Shape Size
et1 ScalarTypeBase Size NoUniqueness
shape1) TypeAliases
t2 =
  Aliases
-> Shape Size -> ScalarTypeBase Size NoUniqueness -> TypeAliases
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array (Aliases
als1 Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> Aliases
aliases TypeAliases
t2) Shape Size
et1 ScalarTypeBase Size NoUniqueness
shape1
combineAliases (Scalar (TypeVar Aliases
als1 QualName VName
tv1 [TypeArg Size]
targs1)) TypeAliases
t2 =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size Aliases
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar (Aliases
als1 Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> Aliases
aliases TypeAliases
t2) QualName VName
tv1 [TypeArg Size]
targs1
combineAliases TypeAliases
t1 (Scalar (TypeVar Aliases
als2 QualName VName
tv2 [TypeArg Size]
targs2)) =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size Aliases
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar (Aliases
als2 Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> Aliases
aliases TypeAliases
t1) QualName VName
tv2 [TypeArg Size]
targs2
combineAliases (Scalar (Record Map Name TypeAliases
ts1)) (Scalar (Record Map Name TypeAliases
ts2))
  | Map Name TypeAliases -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name TypeAliases
ts1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name TypeAliases -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name TypeAliases
ts2,
    [Name] -> [Name]
forall a. Ord a => [a] -> [a]
L.sort (Map Name TypeAliases -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name TypeAliases
ts1) [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name] -> [Name]
forall a. Ord a => [a] -> [a]
L.sort (Map Name TypeAliases -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name TypeAliases
ts2) =
      ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Map Name TypeAliases -> ScalarTypeBase Size Aliases
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name TypeAliases -> ScalarTypeBase Size Aliases)
-> Map Name TypeAliases -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ (TypeAliases -> TypeAliases -> TypeAliases)
-> Map Name TypeAliases
-> Map Name TypeAliases
-> Map Name TypeAliases
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeAliases -> TypeAliases -> TypeAliases
combineAliases Map Name TypeAliases
ts1 Map Name TypeAliases
ts2
combineAliases
  (Scalar (Arrow Aliases
als1 PName
mn1 Diet
d1 StructType
pt1 (RetType [VName]
dims1 ResType
rt1)))
  (Scalar (Arrow Aliases
als2 PName
_ Diet
_ StructType
_ (RetType [VName]
_ ResType
_))) =
    ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Aliases
-> PName
-> Diet
-> StructType
-> RetTypeBase Size Uniqueness
-> ScalarTypeBase Size Aliases
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow (Aliases
als1 Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> Aliases
als2) PName
mn1 Diet
d1 StructType
pt1 ([VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims1 ResType
rt1))
combineAliases (Scalar (Sum Map Name [TypeAliases]
cs1)) (Scalar (Sum Map Name [TypeAliases]
cs2))
  | Map Name [TypeAliases] -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [TypeAliases]
cs1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name [TypeAliases] -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [TypeAliases]
cs2,
    [Name] -> [Name]
forall a. Ord a => [a] -> [a]
L.sort (Map Name [TypeAliases] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeAliases]
cs1) [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name] -> [Name]
forall a. Ord a => [a] -> [a]
L.sort (Map Name [TypeAliases] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeAliases]
cs2) =
      ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Map Name [TypeAliases] -> ScalarTypeBase Size Aliases
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeAliases] -> ScalarTypeBase Size Aliases)
-> Map Name [TypeAliases] -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ ([TypeAliases] -> [TypeAliases] -> [TypeAliases])
-> Map Name [TypeAliases]
-> Map Name [TypeAliases]
-> Map Name [TypeAliases]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((TypeAliases -> TypeAliases -> TypeAliases)
-> [TypeAliases] -> [TypeAliases] -> [TypeAliases]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeAliases -> TypeAliases -> TypeAliases
combineAliases) Map Name [TypeAliases]
cs1 Map Name [TypeAliases]
cs2
combineAliases (Scalar (Prim PrimType
t)) TypeAliases
_ = ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size Aliases
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t
combineAliases TypeAliases
t1 TypeAliases
t2 =
  String -> TypeAliases
forall a. HasCallStack => String -> a
error (String -> TypeAliases) -> String -> TypeAliases
forall a b. (a -> b) -> a -> b
$ String
"combineAliases invalid args: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (TypeAliases, TypeAliases) -> String
forall a. Show a => a -> String
show (TypeAliases
t1, TypeAliases
t2)

-- An alias inhibits uniqueness if it is used in disjoint values.
aliasesMultipleTimes :: TypeAliases -> Names
aliasesMultipleTimes :: TypeAliases -> Names
aliasesMultipleTimes = [VName] -> Names
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Names)
-> (TypeAliases -> [VName]) -> TypeAliases -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Int) -> VName) -> [(VName, Int)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Int) -> VName
forall a b. (a, b) -> a
fst ([(VName, Int)] -> [VName])
-> (TypeAliases -> [(VName, Int)]) -> TypeAliases -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Int) -> Bool) -> [(VName, Int)] -> [(VName, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (Int -> Bool) -> ((VName, Int) -> Int) -> (VName, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Int) -> Int
forall a b. (a, b) -> b
snd) ([(VName, Int)] -> [(VName, Int)])
-> (TypeAliases -> [(VName, Int)]) -> TypeAliases -> [(VName, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Int -> [(VName, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName Int -> [(VName, Int)])
-> (TypeAliases -> Map VName Int) -> TypeAliases -> [(VName, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeAliases -> Map VName Int
delve
  where
    delve :: TypeAliases -> Map VName Int
delve (Scalar (Record Map Name TypeAliases
fs)) =
      (Map VName Int -> Map VName Int -> Map VName Int)
-> Map VName Int -> [Map VName Int] -> Map VName Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Int -> Int -> Int)
-> Map VName Int -> Map VName Int -> Map VName Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)) Map VName Int
forall a. Monoid a => a
mempty ([Map VName Int] -> Map VName Int)
-> [Map VName Int] -> Map VName Int
forall a b. (a -> b) -> a -> b
$ (TypeAliases -> Map VName Int) -> [TypeAliases] -> [Map VName Int]
forall a b. (a -> b) -> [a] -> [b]
map TypeAliases -> Map VName Int
delve ([TypeAliases] -> [Map VName Int])
-> [TypeAliases] -> [Map VName Int]
forall a b. (a -> b) -> a -> b
$ Map Name TypeAliases -> [TypeAliases]
forall k a. Map k a -> [a]
M.elems Map Name TypeAliases
fs
    delve TypeAliases
t =
      [(VName, Int)] -> Map VName Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Int)] -> Map VName Int)
-> [(VName, Int)] -> Map VName Int
forall a b. (a -> b) -> a -> b
$ (Alias -> (VName, Int)) -> [Alias] -> [(VName, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((,Int
1 :: Int) (VName -> (VName, Int))
-> (Alias -> VName) -> Alias -> (VName, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) ([Alias] -> [(VName, Int)]) -> [Alias] -> [(VName, Int)]
forall a b. (a -> b) -> a -> b
$ Aliases -> [Alias]
forall a. Set a -> [a]
S.toList (Aliases -> [Alias]) -> Aliases -> [Alias]
forall a b. (a -> b) -> a -> b
$ TypeAliases -> Aliases
aliases TypeAliases
t

consumingParams :: [Pat ParamType] -> Names
consumingParams :: [Pat (TypeBase Size Diet)] -> Names
consumingParams =
  [VName] -> Names
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Names)
-> ([Pat (TypeBase Size Diet)] -> [VName])
-> [Pat (TypeBase Size Diet)]
-> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, TypeBase Size Diet) -> VName)
-> [(VName, TypeBase Size Diet)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, TypeBase Size Diet) -> VName
forall a b. (a, b) -> a
fst ([(VName, TypeBase Size Diet)] -> [VName])
-> ([Pat (TypeBase Size Diet)] -> [(VName, TypeBase Size Diet)])
-> [Pat (TypeBase Size Diet)]
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, TypeBase Size Diet) -> Bool)
-> [(VName, TypeBase Size Diet)] -> [(VName, TypeBase Size Diet)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume) (Diet -> Bool)
-> ((VName, TypeBase Size Diet) -> Diet)
-> (VName, TypeBase Size Diet)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Size Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet (TypeBase Size Diet -> Diet)
-> ((VName, TypeBase Size Diet) -> TypeBase Size Diet)
-> (VName, TypeBase Size Diet)
-> Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, TypeBase Size Diet) -> TypeBase Size Diet
forall a b. (a, b) -> b
snd) ([(VName, TypeBase Size Diet)] -> [(VName, TypeBase Size Diet)])
-> ([Pat (TypeBase Size Diet)] -> [(VName, TypeBase Size Diet)])
-> [Pat (TypeBase Size Diet)]
-> [(VName, TypeBase Size Diet)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Pat (TypeBase Size Diet) -> [(VName, TypeBase Size Diet)])
-> [Pat (TypeBase Size Diet)] -> [(VName, TypeBase Size Diet)]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat (TypeBase Size Diet) -> [(VName, TypeBase Size Diet)]
forall t. Pat t -> [(VName, t)]
patternMap

arrayAliases :: TypeAliases -> Aliases
arrayAliases :: TypeAliases -> Aliases
arrayAliases (Array Aliases
als Shape Size
_ ScalarTypeBase Size NoUniqueness
_) = Aliases
als
arrayAliases (Scalar Prim {}) = Aliases
forall a. Monoid a => a
mempty
arrayAliases (Scalar (Record Map Name TypeAliases
fs)) = (TypeAliases -> Aliases) -> Map Name TypeAliases -> Aliases
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeAliases -> Aliases
arrayAliases Map Name TypeAliases
fs
arrayAliases (Scalar (TypeVar Aliases
als QualName VName
_ [TypeArg Size]
_)) = Aliases
als
arrayAliases (Scalar Arrow {}) = Aliases
forall a. Monoid a => a
mempty
arrayAliases (Scalar (Sum Map Name [TypeAliases]
fs)) =
  [Aliases] -> Aliases
forall a. Monoid a => [a] -> a
mconcat ([Aliases] -> Aliases) -> [Aliases] -> Aliases
forall a b. (a -> b) -> a -> b
$ ([TypeAliases] -> [Aliases]) -> [[TypeAliases]] -> [Aliases]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((TypeAliases -> Aliases) -> [TypeAliases] -> [Aliases]
forall a b. (a -> b) -> [a] -> [b]
map TypeAliases -> Aliases
arrayAliases) ([[TypeAliases]] -> [Aliases]) -> [[TypeAliases]] -> [Aliases]
forall a b. (a -> b) -> a -> b
$ Map Name [TypeAliases] -> [[TypeAliases]]
forall k a. Map k a -> [a]
M.elems Map Name [TypeAliases]
fs

overlapCheck :: (Pretty src, Pretty ve) => Loc -> (src, TypeAliases) -> (ve, TypeAliases) -> CheckM ()
overlapCheck :: forall src ve.
(Pretty src, Pretty ve) =>
Loc -> (src, TypeAliases) -> (ve, TypeAliases) -> CheckM ()
overlapCheck Loc
loc (src
src, TypeAliases
src_als) (ve
ve, TypeAliases
ve_als) =
  Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((Alias -> Bool) -> Aliases -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Alias -> Aliases -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TypeAliases -> Aliases
aliases TypeAliases
src_als) (TypeAliases -> Aliases
aliases TypeAliases
ve_als)) (CheckM () -> CheckM ()) -> CheckM () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
    Loc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
      Doc ()
"Source array for in-place update"
        Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (src -> Doc ()
forall ann. src -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty src
src)
        Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"might alias update value"
        Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (ve -> Doc ()
forall ann. ve -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ve
ve)
        Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"Hint: use"
        Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes Doc ()
"copy"
        Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
"to remove aliases from the value."

inferReturnUniqueness :: [Pat ParamType] -> ResType -> TypeAliases -> ResType
inferReturnUniqueness :: [Pat (TypeBase Size Diet)] -> ResType -> TypeAliases -> ResType
inferReturnUniqueness [] ResType
ret TypeAliases
_ = ResType
ret ResType -> Uniqueness -> ResType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setUniqueness` Uniqueness
Nonunique
inferReturnUniqueness [Pat (TypeBase Size Diet)]
params ResType
ret TypeAliases
ret_als = ResType -> TypeAliases -> ResType
forall {dim} {u1}.
TypeBase dim u1 -> TypeAliases -> TypeBase dim Uniqueness
delve ResType
ret TypeAliases
ret_als
  where
    forbidden :: Names
forbidden = TypeAliases -> Names
aliasesMultipleTimes TypeAliases
ret_als
    consumings :: Names
consumings = [Pat (TypeBase Size Diet)] -> Names
consumingParams [Pat (TypeBase Size Diet)]
params
    delve :: TypeBase dim u1 -> TypeAliases -> TypeBase dim Uniqueness
delve (Scalar (Record Map Name (TypeBase dim u1)
fs1)) (Scalar (Record Map Name TypeAliases
fs2)) =
      ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness)
-> ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim Uniqueness) -> ScalarTypeBase dim Uniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase dim Uniqueness)
 -> ScalarTypeBase dim Uniqueness)
-> Map Name (TypeBase dim Uniqueness)
-> ScalarTypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ (TypeBase dim u1 -> TypeAliases -> TypeBase dim Uniqueness)
-> Map Name (TypeBase dim u1)
-> Map Name TypeAliases
-> Map Name (TypeBase dim Uniqueness)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase dim u1 -> TypeAliases -> TypeBase dim Uniqueness
delve Map Name (TypeBase dim u1)
fs1 Map Name TypeAliases
fs2
    delve (Scalar (Sum Map Name [TypeBase dim u1]
cs1)) (Scalar (Sum Map Name [TypeAliases]
cs2)) =
      ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness)
-> ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase dim Uniqueness] -> ScalarTypeBase dim Uniqueness
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase dim Uniqueness]
 -> ScalarTypeBase dim Uniqueness)
-> Map Name [TypeBase dim Uniqueness]
-> ScalarTypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim u1] -> [TypeAliases] -> [TypeBase dim Uniqueness])
-> Map Name [TypeBase dim u1]
-> Map Name [TypeAliases]
-> Map Name [TypeBase dim Uniqueness]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((TypeBase dim u1 -> TypeAliases -> TypeBase dim Uniqueness)
-> [TypeBase dim u1] -> [TypeAliases] -> [TypeBase dim Uniqueness]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase dim u1 -> TypeAliases -> TypeBase dim Uniqueness
delve) Map Name [TypeBase dim u1]
cs1 Map Name [TypeAliases]
cs2
    delve TypeBase dim u1
t TypeAliases
t_als
      | (VName -> Bool) -> Names -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (VName -> Names -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Names
consumings) (Names -> Bool) -> Names -> Bool
forall a b. (a -> b) -> a -> b
$ Aliases -> Names
boundAliases (TypeAliases -> Aliases
arrayAliases TypeAliases
t_als),
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Alias -> Bool) -> Aliases -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> Names -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Names
forbidden) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (TypeAliases -> Aliases
aliases TypeAliases
t_als) =
          TypeBase dim u1
t TypeBase dim u1 -> Uniqueness -> TypeBase dim Uniqueness
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setUniqueness` Uniqueness
Unique
      | Bool
otherwise =
          TypeBase dim u1
t TypeBase dim u1 -> Uniqueness -> TypeBase dim Uniqueness
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setUniqueness` Uniqueness
Nonunique

checkSubExps :: (ASTMappable e) => e -> CheckM e
checkSubExps :: forall e. ASTMappable e => e -> CheckM e
checkSubExps = ASTMapper CheckM -> e -> CheckM e
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> e -> m e
astMap ASTMapper CheckM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = fmap fst . checkExp}

noAliases :: Exp -> CheckM (Exp, TypeAliases)
noAliases :: Size -> CheckM (Size, TypeAliases)
noAliases Size
e = do
  e' <- Size -> CheckM Size
forall e. ASTMappable e => e -> CheckM e
checkSubExps Size
e
  pure (e', second (const mempty) (typeOf e))

aliasParts :: TypeAliases -> [Aliases]
aliasParts :: TypeAliases -> [Aliases]
aliasParts (Scalar (Record Map Name TypeAliases
ts)) = (TypeAliases -> [Aliases]) -> [TypeAliases] -> [Aliases]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeAliases -> [Aliases]
aliasParts ([TypeAliases] -> [Aliases]) -> [TypeAliases] -> [Aliases]
forall a b. (a -> b) -> a -> b
$ Map Name TypeAliases -> [TypeAliases]
forall k a. Map k a -> [a]
M.elems Map Name TypeAliases
ts
aliasParts TypeAliases
t = [TypeAliases -> Aliases
aliases TypeAliases
t]

noSelfAliases :: Loc -> TypeAliases -> CheckM ()
noSelfAliases :: Loc -> TypeAliases -> CheckM ()
noSelfAliases Loc
loc = (Aliases -> Aliases -> CheckM Aliases)
-> Aliases -> [Aliases] -> CheckM ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ Aliases -> Aliases -> CheckM Aliases
forall {a}. Ord a => Set a -> Set a -> CheckM (Set a)
check Aliases
forall a. Monoid a => a
mempty ([Aliases] -> CheckM ())
-> (TypeAliases -> [Aliases]) -> TypeAliases -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeAliases -> [Aliases]
aliasParts
  where
    check :: Set a -> Set a -> CheckM (Set a)
check Set a
seen Set a
als = do
      Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((a -> Bool) -> Set a -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set a
seen) Set a
als) (CheckM () -> CheckM ()) -> CheckM () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
        Loc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError Loc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> CheckM ()) -> (Doc () -> Doc ()) -> Doc () -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
withIndexLink Doc ()
"self-aliases-arg" (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
          Doc ()
"Argument passed for consuming parameter is self-aliased."
      Set a -> CheckM (Set a)
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set a -> CheckM (Set a)) -> Set a -> CheckM (Set a)
forall a b. (a -> b) -> a -> b
$ Set a
als Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> Set a
seen

consumeAsNeeded :: Loc -> ParamType -> TypeAliases -> CheckM ()
consumeAsNeeded :: Loc -> TypeBase Size Diet -> TypeAliases -> CheckM ()
consumeAsNeeded Loc
loc (Scalar (Record Map Name (TypeBase Size Diet)
fs1)) (Scalar (Record Map Name TypeAliases
fs2)) =
  [CheckM ()] -> CheckM ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CheckM ()] -> CheckM ()) -> [CheckM ()] -> CheckM ()
forall a b. (a -> b) -> a -> b
$ Map Name (CheckM ()) -> [CheckM ()]
forall k a. Map k a -> [a]
M.elems (Map Name (CheckM ()) -> [CheckM ()])
-> Map Name (CheckM ()) -> [CheckM ()]
forall a b. (a -> b) -> a -> b
$ (TypeBase Size Diet -> TypeAliases -> CheckM ())
-> Map Name (TypeBase Size Diet)
-> Map Name TypeAliases
-> Map Name (CheckM ())
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (Loc -> TypeBase Size Diet -> TypeAliases -> CheckM ()
consumeAsNeeded Loc
loc) Map Name (TypeBase Size Diet)
fs1 Map Name TypeAliases
fs2
consumeAsNeeded Loc
loc TypeBase Size Diet
pt TypeAliases
t =
  Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (TypeBase Size Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase Size Diet
pt Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume) (CheckM () -> CheckM ()) -> CheckM () -> CheckM ()
forall a b. (a -> b) -> a -> b
$ Loc -> Aliases -> CheckM ()
consumeAliases Loc
loc (Aliases -> CheckM ()) -> Aliases -> CheckM ()
forall a b. (a -> b) -> a -> b
$ TypeAliases -> Aliases
aliases TypeAliases
t

checkArg :: [(Exp, TypeAliases)] -> ParamType -> Exp -> CheckM (Exp, TypeAliases)
checkArg :: [(Size, TypeAliases)]
-> TypeBase Size Diet -> Size -> CheckM (Size, TypeAliases)
checkArg [(Size, TypeAliases)]
prev TypeBase Size Diet
p_t Size
e = do
  ((e', e_als), e_cons) <- CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (Size, TypeAliases)
 -> CheckM ((Size, TypeAliases), Consumed))
-> CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  consumed e_cons
  let e_t = Size -> StructType
typeOf Size
e'
  when (e_cons /= mempty && not (orderZero e_t)) $
    addError (locOf e) mempty $
      "Argument of functional type"
        </> indent 2 (pretty e_t)
        </> "contains consumption, which is not allowed."
  when (diet p_t == Consume) $ do
    noSelfAliases (locOf e) e_als
    consumeAsNeeded (locOf e) p_t e_als
    case mapMaybe prevAlias $ S.toList $ boundAliases $ aliases e_als of
      [] -> () -> CheckM ()
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      (VName
v, Size
prev_arg) : [(VName, Size)]
_ ->
        Loc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError (Size -> Loc
forall a. Located a => a -> Loc
locOf Size
e) Notes
forall a. Monoid a => a
mempty (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
          Doc ()
"Argument is consumed, but aliases"
            Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v)
            Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"which is also aliased by other argument"
            Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Size -> Doc ()
forall a ann. Pretty a => a -> Doc ann
forall ann. Size -> Doc ann
pretty Size
prev_arg)
            Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"at"
            Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Text -> Doc ()
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Loc -> Loc -> Text
forall a b. (Located a, Located b) => a -> b -> Text
locTextRel (Size -> Loc
forall a. Located a => a -> Loc
locOf Size
e) (Size -> Loc
forall a. Located a => a -> Loc
locOf Size
prev_arg))
            Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
  pure (e', e_als)
  where
    prevAlias :: VName -> Maybe (VName, Size)
prevAlias VName
v =
      (VName
v,) (Size -> (VName, Size))
-> ((Size, TypeAliases) -> Size)
-> (Size, TypeAliases)
-> (VName, Size)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Size, TypeAliases) -> Size
forall a b. (a, b) -> a
fst ((Size, TypeAliases) -> (VName, Size))
-> Maybe (Size, TypeAliases) -> Maybe (VName, Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Size, TypeAliases) -> Bool)
-> [(Size, TypeAliases)] -> Maybe (Size, TypeAliases)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (VName -> Names -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member VName
v (Names -> Bool)
-> ((Size, TypeAliases) -> Names) -> (Size, TypeAliases) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliases -> Names
boundAliases (Aliases -> Names)
-> ((Size, TypeAliases) -> Aliases) -> (Size, TypeAliases) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeAliases -> Aliases
aliases (TypeAliases -> Aliases)
-> ((Size, TypeAliases) -> TypeAliases)
-> (Size, TypeAliases)
-> Aliases
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Size, TypeAliases) -> TypeAliases
forall a b. (a, b) -> b
snd) [(Size, TypeAliases)]
prev

-- | @returnType appres ret_type arg_diet arg_type@ gives result of applying
-- an argument the given types to a function with the given return
-- type, consuming the argument with the given diet.
returnType :: Aliases -> ResType -> Diet -> TypeAliases -> TypeAliases
returnType :: Aliases -> ResType -> Diet -> TypeAliases -> TypeAliases
returnType Aliases
_ (Array Uniqueness
Unique Shape Size
et ScalarTypeBase Size NoUniqueness
shape) Diet
_ TypeAliases
_ =
  Aliases
-> Shape Size -> ScalarTypeBase Size NoUniqueness -> TypeAliases
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Aliases
forall a. Monoid a => a
mempty Shape Size
et ScalarTypeBase Size NoUniqueness
shape
returnType Aliases
appres (Array Uniqueness
Nonunique Shape Size
et ScalarTypeBase Size NoUniqueness
shape) Diet
Consume TypeAliases
_ =
  Aliases
-> Shape Size -> ScalarTypeBase Size NoUniqueness -> TypeAliases
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Aliases
appres Shape Size
et ScalarTypeBase Size NoUniqueness
shape
returnType Aliases
appres (Array Uniqueness
Nonunique Shape Size
et ScalarTypeBase Size NoUniqueness
shape) Diet
Observe TypeAliases
arg =
  Aliases
-> Shape Size -> ScalarTypeBase Size NoUniqueness -> TypeAliases
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array (Aliases
appres Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> Aliases
aliases TypeAliases
arg) Shape Size
et ScalarTypeBase Size NoUniqueness
shape
returnType Aliases
_ (Scalar (TypeVar Uniqueness
Unique QualName VName
t [TypeArg Size]
targs)) Diet
_ TypeAliases
_ =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size Aliases
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar Aliases
forall a. Monoid a => a
mempty QualName VName
t [TypeArg Size]
targs
returnType Aliases
appres (Scalar (TypeVar Uniqueness
Nonunique QualName VName
t [TypeArg Size]
targs)) Diet
Consume TypeAliases
_ =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size Aliases
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar Aliases
appres QualName VName
t [TypeArg Size]
targs
returnType Aliases
appres (Scalar (TypeVar Uniqueness
Nonunique QualName VName
t [TypeArg Size]
targs)) Diet
Observe TypeAliases
arg =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> QualName VName -> [TypeArg Size] -> ScalarTypeBase Size Aliases
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar (Aliases
appres Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> Aliases
aliases TypeAliases
arg) QualName VName
t [TypeArg Size]
targs
returnType Aliases
appres (Scalar (Record Map Name ResType
fs)) Diet
d TypeAliases
arg =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Map Name TypeAliases -> ScalarTypeBase Size Aliases
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name TypeAliases -> ScalarTypeBase Size Aliases)
-> Map Name TypeAliases -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ (ResType -> TypeAliases)
-> Map Name ResType -> Map Name TypeAliases
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ResType
et -> Aliases -> ResType -> Diet -> TypeAliases -> TypeAliases
returnType Aliases
appres ResType
et Diet
d TypeAliases
arg) Map Name ResType
fs
returnType Aliases
_ (Scalar (Prim PrimType
t)) Diet
_ TypeAliases
_ =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size Aliases
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t
returnType Aliases
appres (Scalar (Arrow Uniqueness
_ PName
v Diet
pd StructType
t1 (RetType [VName]
dims ResType
t2))) Diet
Consume TypeAliases
_ =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> PName
-> Diet
-> StructType
-> RetTypeBase Size Uniqueness
-> ScalarTypeBase Size Aliases
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow Aliases
appres PName
v Diet
pd StructType
t1 (RetTypeBase Size Uniqueness -> ScalarTypeBase Size Aliases)
-> RetTypeBase Size Uniqueness -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims ResType
t2
returnType Aliases
appres (Scalar (Arrow Uniqueness
_ PName
v Diet
pd StructType
t1 (RetType [VName]
dims ResType
t2))) Diet
Observe TypeAliases
arg =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Aliases
-> PName
-> Diet
-> StructType
-> RetTypeBase Size Uniqueness
-> ScalarTypeBase Size Aliases
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow (Aliases
appres Aliases -> Aliases -> Aliases
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> Aliases
aliases TypeAliases
arg) PName
v Diet
pd StructType
t1 (RetTypeBase Size Uniqueness -> ScalarTypeBase Size Aliases)
-> RetTypeBase Size Uniqueness -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims ResType
t2
returnType Aliases
appres (Scalar (Sum Map Name [ResType]
cs)) Diet
d TypeAliases
arg =
  ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> ScalarTypeBase Size Aliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ Map Name [TypeAliases] -> ScalarTypeBase Size Aliases
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeAliases] -> ScalarTypeBase Size Aliases)
-> Map Name [TypeAliases] -> ScalarTypeBase Size Aliases
forall a b. (a -> b) -> a -> b
$ (([ResType] -> [TypeAliases])
-> Map Name [ResType] -> Map Name [TypeAliases]
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([ResType] -> [TypeAliases])
 -> Map Name [ResType] -> Map Name [TypeAliases])
-> ((ResType -> TypeAliases) -> [ResType] -> [TypeAliases])
-> (ResType -> TypeAliases)
-> Map Name [ResType]
-> Map Name [TypeAliases]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ResType -> TypeAliases) -> [ResType] -> [TypeAliases]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (\ResType
et -> Aliases -> ResType -> Diet -> TypeAliases -> TypeAliases
returnType Aliases
appres ResType
et Diet
d TypeAliases
arg) Map Name [ResType]
cs

applyArg :: TypeAliases -> TypeAliases -> TypeAliases
applyArg :: TypeAliases -> TypeAliases -> TypeAliases
applyArg (Scalar (Arrow Aliases
closure_als PName
_ Diet
d StructType
_ (RetType [VName]
_ ResType
rettype))) TypeAliases
arg_als =
  Aliases -> ResType -> Diet -> TypeAliases -> TypeAliases
returnType Aliases
closure_als ResType
rettype Diet
d TypeAliases
arg_als
applyArg TypeAliases
t TypeAliases
_ = String -> TypeAliases
forall a. HasCallStack => String -> a
error (String -> TypeAliases) -> String -> TypeAliases
forall a b. (a -> b) -> a -> b
$ String
"applyArg: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> String
forall a. Show a => a -> String
show TypeAliases
t

boundFreeInExp :: Exp -> CheckM (M.Map VName TypeAliases)
boundFreeInExp :: Size -> CheckM (Map VName TypeAliases)
boundFreeInExp Size
e = do
  vtable <- (CheckEnv -> Map VName (Entry TypeAliases))
-> CheckM (Map VName (Entry TypeAliases))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CheckEnv -> Map VName (Entry TypeAliases)
envVtable
  pure $
    M.mapMaybe (fmap entryAliases) . M.fromSet (`M.lookup` vtable) $
      fvVars (freeInExp e)

-- Loops are tricky because we want to infer the uniqueness of their
-- parameters.  This is pretty unusual: we do not do this for ordinary
-- functions.
type Loop = (Pat ParamType, LoopInitBase Info VName, LoopFormBase Info VName, Exp)

-- | Mark bindings of consumed names as Consume, except those under a
-- 'PatAscription', which are left unchanged.
updateParamDiet :: (VName -> Bool) -> Pat ParamType -> Pat ParamType
updateParamDiet :: (VName -> Bool)
-> Pat (TypeBase Size Diet) -> Pat (TypeBase Size Diet)
updateParamDiet VName -> Bool
cons = Pat (TypeBase Size Diet) -> Pat (TypeBase Size Diet)
forall {dim}.
PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse
  where
    recurse :: PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse (Wildcard (Info TypeBase dim Diet
t) SrcLoc
wloc) =
      Info (TypeBase dim Diet)
-> SrcLoc -> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (TypeBase dim Diet -> Info (TypeBase dim Diet)
forall a. a -> Info a
Info (TypeBase dim Diet -> Info (TypeBase dim Diet))
-> TypeBase dim Diet -> Info (TypeBase dim Diet)
forall a b. (a -> b) -> a -> b
$ TypeBase dim Diet
t TypeBase dim Diet -> Diet -> TypeBase dim Diet
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setUniqueness` Diet
Observe) SrcLoc
wloc
    recurse (PatParens PatBase Info VName (TypeBase dim Diet)
p SrcLoc
ploc) =
      PatBase Info VName (TypeBase dim Diet)
-> SrcLoc -> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t.
PatBase f vn t -> SrcLoc -> PatBase f vn t
PatParens (PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse PatBase Info VName (TypeBase dim Diet)
p) SrcLoc
ploc
    recurse (PatAttr AttrInfo VName
attr PatBase Info VName (TypeBase dim Diet)
p SrcLoc
ploc) =
      AttrInfo VName
-> PatBase Info VName (TypeBase dim Diet)
-> SrcLoc
-> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t.
AttrInfo vn -> PatBase f vn t -> SrcLoc -> PatBase f vn t
PatAttr AttrInfo VName
attr (PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse PatBase Info VName (TypeBase dim Diet)
p) SrcLoc
ploc
    recurse (Id VName
name (Info TypeBase dim Diet
t) SrcLoc
iloc)
      | VName -> Bool
cons VName
name =
          let t' :: TypeBase dim Diet
t' = TypeBase dim Diet
t TypeBase dim Diet -> Diet -> TypeBase dim Diet
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setUniqueness` Diet
Consume
           in VName
-> Info (TypeBase dim Diet)
-> SrcLoc
-> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
name (TypeBase dim Diet -> Info (TypeBase dim Diet)
forall a. a -> Info a
Info TypeBase dim Diet
t') SrcLoc
iloc
      | Bool
otherwise =
          let t' :: TypeBase dim Diet
t' = TypeBase dim Diet
t TypeBase dim Diet -> Diet -> TypeBase dim Diet
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setUniqueness` Diet
Observe
           in VName
-> Info (TypeBase dim Diet)
-> SrcLoc
-> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
name (TypeBase dim Diet -> Info (TypeBase dim Diet)
forall a. a -> Info a
Info TypeBase dim Diet
t') SrcLoc
iloc
    recurse (TuplePat [PatBase Info VName (TypeBase dim Diet)]
pats SrcLoc
ploc) =
      [PatBase Info VName (TypeBase dim Diet)]
-> SrcLoc -> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat ((PatBase Info VName (TypeBase dim Diet)
 -> PatBase Info VName (TypeBase dim Diet))
-> [PatBase Info VName (TypeBase dim Diet)]
-> [PatBase Info VName (TypeBase dim Diet)]
forall a b. (a -> b) -> [a] -> [b]
map PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse [PatBase Info VName (TypeBase dim Diet)]
pats) SrcLoc
ploc
    recurse (RecordPat [(L Name, PatBase Info VName (TypeBase dim Diet))]
fs SrcLoc
ploc) =
      [(L Name, PatBase Info VName (TypeBase dim Diet))]
-> SrcLoc -> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat (((L Name, PatBase Info VName (TypeBase dim Diet))
 -> (L Name, PatBase Info VName (TypeBase dim Diet)))
-> [(L Name, PatBase Info VName (TypeBase dim Diet))]
-> [(L Name, PatBase Info VName (TypeBase dim Diet))]
forall a b. (a -> b) -> [a] -> [b]
map ((PatBase Info VName (TypeBase dim Diet)
 -> PatBase Info VName (TypeBase dim Diet))
-> (L Name, PatBase Info VName (TypeBase dim Diet))
-> (L Name, PatBase Info VName (TypeBase dim Diet))
forall a b. (a -> b) -> (L Name, a) -> (L Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse) [(L Name, PatBase Info VName (TypeBase dim Diet))]
fs) SrcLoc
ploc
    recurse (PatAscription PatBase Info VName (TypeBase dim Diet)
p TypeExp Size VName
t SrcLoc
ploc) =
      PatBase Info VName (TypeBase dim Diet)
-> TypeExp Size VName
-> SrcLoc
-> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t.
PatBase f vn t
-> TypeExp (ExpBase f vn) vn -> SrcLoc -> PatBase f vn t
PatAscription PatBase Info VName (TypeBase dim Diet)
p TypeExp Size VName
t SrcLoc
ploc
    recurse p :: PatBase Info VName (TypeBase dim Diet)
p@PatLit {} = PatBase Info VName (TypeBase dim Diet)
p
    recurse (PatConstr Name
n Info (TypeBase dim Diet)
t [PatBase Info VName (TypeBase dim Diet)]
ps SrcLoc
ploc) =
      Name
-> Info (TypeBase dim Diet)
-> [PatBase Info VName (TypeBase dim Diet)]
-> SrcLoc
-> PatBase Info VName (TypeBase dim Diet)
forall (f :: * -> *) vn t.
Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
PatConstr Name
n Info (TypeBase dim Diet)
t ((PatBase Info VName (TypeBase dim Diet)
 -> PatBase Info VName (TypeBase dim Diet))
-> [PatBase Info VName (TypeBase dim Diet)]
-> [PatBase Info VName (TypeBase dim Diet)]
forall a b. (a -> b) -> [a] -> [b]
map PatBase Info VName (TypeBase dim Diet)
-> PatBase Info VName (TypeBase dim Diet)
recurse [PatBase Info VName (TypeBase dim Diet)]
ps) SrcLoc
ploc

convergeLoopParam :: Loc -> Pat ParamType -> Names -> TypeAliases -> CheckM (Pat ParamType)
convergeLoopParam :: Loc
-> Pat (TypeBase Size Diet)
-> Names
-> TypeAliases
-> CheckM (Pat (TypeBase Size Diet))
convergeLoopParam Loc
loop_loc Pat (TypeBase Size Diet)
param Names
body_cons TypeAliases
body_als = do
  let -- Make the pattern Consume where needed.
      param' :: Pat (TypeBase Size Diet)
param' = (VName -> Bool)
-> Pat (TypeBase Size Diet) -> Pat (TypeBase Size Diet)
updateParamDiet (VName -> Names -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` (VName -> Bool) -> Names -> Names
forall a. (a -> Bool) -> Set a -> Set a
S.filter (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Pat (TypeBase Size Diet) -> [VName]
forall t. Pat t -> [VName]
patNames Pat (TypeBase Size Diet)
param) Names
body_cons) Pat (TypeBase Size Diet)
param

  -- Check that the new values of consumed merge parameters do not
  -- alias something bound outside the loop, AND that anything
  -- returned for a unique merge parameter does not alias anything
  -- else returned.
  let checkMergeReturn :: PatBase Info vn (TypeBase shape Diet)
-> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet))
checkMergeReturn (Id vn
pat_v (Info TypeBase shape Diet
pat_v_t) SrcLoc
patloc) TypeAliases
t = do
        let free_als :: Names
free_als = (VName -> Bool) -> Names -> Names
forall a. (a -> Bool) -> Set a -> Set a
S.filter (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Pat (TypeBase Size Diet) -> [VName]
forall t. Pat t -> [VName]
patNames Pat (TypeBase Size Diet)
param) (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$ Aliases -> Names
boundAliases (TypeAliases -> Aliases
aliases TypeAliases
t)
        Bool -> t CheckM () -> t CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (TypeBase shape Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase shape Diet
pat_v_t Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume) (t CheckM () -> t CheckM ()) -> t CheckM () -> t CheckM ()
forall a b. (a -> b) -> a -> b
$ Names -> (VName -> t CheckM ()) -> t CheckM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Names
free_als ((VName -> t CheckM ()) -> t CheckM ())
-> (VName -> t CheckM ()) -> t CheckM ()
forall a b. (a -> b) -> a -> b
$ \VName
v ->
          CheckM () -> t CheckM ()
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CheckM () -> t CheckM ())
-> (Doc () -> CheckM ()) -> Doc () -> t CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Loc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError Loc
loop_loc Notes
forall a. Monoid a => a
mempty (Doc () -> t CheckM ()) -> Doc () -> t CheckM ()
forall a b. (a -> b) -> a -> b
$
            Doc ()
"Return value for consuming loop parameter"
              Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (vn -> Doc ()
forall a. vn -> Doc a
forall v a. IsName v => v -> Doc a
prettyName vn
pat_v)
              Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
"aliases"
              Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v)
              Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
        (cons, obs) <- t CheckM (Aliases, Aliases)
forall s (m :: * -> *). MonadState s m => m s
get
        unless (S.null $ aliases t `S.intersection` cons) $
          lift . addError loop_loc mempty $
            "Return value for loop parameter"
              <+> dquotes (prettyName pat_v)
              <+> "aliases other consumed loop parameter."
        when
          ( diet pat_v_t == Consume
              && not (S.null (aliases t `S.intersection` (cons <> obs)))
          )
          $ lift . addError loop_loc mempty
          $ withIndexLink "aliases-previously-returned"
          $ "Return value for consuming loop parameter"
            <+> dquotes (prettyName pat_v)
            <+> "aliases previously returned value."
        if diet pat_v_t == Consume
          then put (cons <> aliases t, obs)
          else put (cons, obs <> aliases t)

        pure $ Id pat_v (Info pat_v_t) patloc
      checkMergeReturn (Wildcard (Info TypeBase shape Diet
pat_v_t) SrcLoc
patloc) TypeAliases
_ =
        PatBase Info vn (TypeBase shape Diet)
-> t CheckM (PatBase Info vn (TypeBase shape Diet))
forall a. a -> t CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatBase Info vn (TypeBase shape Diet)
 -> t CheckM (PatBase Info vn (TypeBase shape Diet)))
-> PatBase Info vn (TypeBase shape Diet)
-> t CheckM (PatBase Info vn (TypeBase shape Diet))
forall a b. (a -> b) -> a -> b
$ Info (TypeBase shape Diet)
-> SrcLoc -> PatBase Info vn (TypeBase shape Diet)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (TypeBase shape Diet -> Info (TypeBase shape Diet)
forall a. a -> Info a
Info TypeBase shape Diet
pat_v_t) SrcLoc
patloc
      checkMergeReturn (PatParens PatBase Info vn (TypeBase shape Diet)
p SrcLoc
_) TypeAliases
t =
        PatBase Info vn (TypeBase shape Diet)
-> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet))
checkMergeReturn PatBase Info vn (TypeBase shape Diet)
p TypeAliases
t
      checkMergeReturn (PatAscription PatBase Info vn (TypeBase shape Diet)
p TypeExp (ExpBase Info vn) vn
_ SrcLoc
_) TypeAliases
t =
        PatBase Info vn (TypeBase shape Diet)
-> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet))
checkMergeReturn PatBase Info vn (TypeBase shape Diet)
p TypeAliases
t
      checkMergeReturn (RecordPat [(L Name, PatBase Info vn (TypeBase shape Diet))]
pfs SrcLoc
patloc) (Scalar (Record Map Name TypeAliases
tfs)) =
        [(L Name, PatBase Info vn (TypeBase shape Diet))]
-> SrcLoc -> PatBase Info vn (TypeBase shape Diet)
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([(L Name, PatBase Info vn (TypeBase shape Diet))]
 -> SrcLoc -> PatBase Info vn (TypeBase shape Diet))
-> (Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
    -> [(L Name, PatBase Info vn (TypeBase shape Diet))])
-> Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
-> SrcLoc
-> PatBase Info vn (TypeBase shape Diet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, (Loc, PatBase Info vn (TypeBase shape Diet)))
 -> (L Name, PatBase Info vn (TypeBase shape Diet)))
-> [(Name, (Loc, PatBase Info vn (TypeBase shape Diet)))]
-> [(L Name, PatBase Info vn (TypeBase shape Diet))]
forall a b. (a -> b) -> [a] -> [b]
map (Name, (Loc, PatBase Info vn (TypeBase shape Diet)))
-> (L Name, PatBase Info vn (TypeBase shape Diet))
forall {a} {b}. (a, (Loc, b)) -> (L a, b)
unshuffle ([(Name, (Loc, PatBase Info vn (TypeBase shape Diet)))]
 -> [(L Name, PatBase Info vn (TypeBase shape Diet))])
-> (Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
    -> [(Name, (Loc, PatBase Info vn (TypeBase shape Diet)))])
-> Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
-> [(L Name, PatBase Info vn (TypeBase shape Diet))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
-> [(Name, (Loc, PatBase Info vn (TypeBase shape Diet)))]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
 -> SrcLoc -> PatBase Info vn (TypeBase shape Diet))
-> t CheckM (Map Name (Loc, PatBase Info vn (TypeBase shape Diet)))
-> t CheckM (SrcLoc -> PatBase Info vn (TypeBase shape Diet))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (t CheckM (Loc, PatBase Info vn (TypeBase shape Diet)))
-> t CheckM (Map Name (Loc, PatBase Info vn (TypeBase shape Diet)))
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Map Name (m a) -> m (Map Name a)
sequence Map Name (t CheckM (Loc, PatBase Info vn (TypeBase shape Diet)))
pfs' t CheckM (SrcLoc -> PatBase Info vn (TypeBase shape Diet))
-> t CheckM SrcLoc
-> t CheckM (PatBase Info vn (TypeBase shape Diet))
forall a b. t CheckM (a -> b) -> t CheckM a -> t CheckM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t CheckM SrcLoc
forall a. a -> t CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
patloc
        where
          pfs' :: Map Name (t CheckM (Loc, PatBase Info vn (TypeBase shape Diet)))
pfs' = ((Loc, PatBase Info vn (TypeBase shape Diet))
 -> TypeAliases
 -> t CheckM (Loc, PatBase Info vn (TypeBase shape Diet)))
-> Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
-> Map Name TypeAliases
-> Map Name (t CheckM (Loc, PatBase Info vn (TypeBase shape Diet)))
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (Loc, PatBase Info vn (TypeBase shape Diet))
-> TypeAliases
-> t CheckM (Loc, PatBase Info vn (TypeBase shape Diet))
forall {t}.
(t, PatBase Info vn (TypeBase shape Diet))
-> TypeAliases
-> t CheckM (t, PatBase Info vn (TypeBase shape Diet))
check ([(Name, (Loc, PatBase Info vn (TypeBase shape Diet)))]
-> Map Name (Loc, PatBase Info vn (TypeBase shape Diet))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (((L Name, PatBase Info vn (TypeBase shape Diet))
 -> (Name, (Loc, PatBase Info vn (TypeBase shape Diet))))
-> [(L Name, PatBase Info vn (TypeBase shape Diet))]
-> [(Name, (Loc, PatBase Info vn (TypeBase shape Diet)))]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, PatBase Info vn (TypeBase shape Diet))
-> (Name, (Loc, PatBase Info vn (TypeBase shape Diet)))
forall {a} {b}. (L a, b) -> (a, (Loc, b))
shuffle [(L Name, PatBase Info vn (TypeBase shape Diet))]
pfs)) Map Name TypeAliases
tfs
          check :: (t, PatBase Info vn (TypeBase shape Diet))
-> TypeAliases
-> t CheckM (t, PatBase Info vn (TypeBase shape Diet))
check (t
loc, PatBase Info vn (TypeBase shape Diet)
x) TypeAliases
y = (t
loc,) (PatBase Info vn (TypeBase shape Diet)
 -> (t, PatBase Info vn (TypeBase shape Diet)))
-> t CheckM (PatBase Info vn (TypeBase shape Diet))
-> t CheckM (t, PatBase Info vn (TypeBase shape Diet))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info vn (TypeBase shape Diet)
-> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet))
checkMergeReturn PatBase Info vn (TypeBase shape Diet)
x TypeAliases
y
          shuffle :: (L a, b) -> (a, (Loc, b))
shuffle (L Loc
loc a
v, b
t) = (a
v, (Loc
loc, b
t))
          unshuffle :: (a, (Loc, b)) -> (L a, b)
unshuffle (a
v, (Loc
loc, b
t)) = (Loc -> a -> L a
forall a. Loc -> a -> L a
L Loc
loc a
v, b
t)
      checkMergeReturn (TuplePat [PatBase Info vn (TypeBase shape Diet)]
pats SrcLoc
patloc) TypeAliases
t
        | Just [TypeAliases]
ts <- TypeAliases -> Maybe [TypeAliases]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord TypeAliases
t =
            [PatBase Info vn (TypeBase shape Diet)]
-> SrcLoc -> PatBase Info vn (TypeBase shape Diet)
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat ([PatBase Info vn (TypeBase shape Diet)]
 -> SrcLoc -> PatBase Info vn (TypeBase shape Diet))
-> t CheckM [PatBase Info vn (TypeBase shape Diet)]
-> t CheckM (SrcLoc -> PatBase Info vn (TypeBase shape Diet))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatBase Info vn (TypeBase shape Diet)
 -> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet)))
-> [PatBase Info vn (TypeBase shape Diet)]
-> [TypeAliases]
-> t CheckM [PatBase Info vn (TypeBase shape Diet)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM PatBase Info vn (TypeBase shape Diet)
-> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet))
checkMergeReturn [PatBase Info vn (TypeBase shape Diet)]
pats [TypeAliases]
ts t CheckM (SrcLoc -> PatBase Info vn (TypeBase shape Diet))
-> t CheckM SrcLoc
-> t CheckM (PatBase Info vn (TypeBase shape Diet))
forall a b. t CheckM (a -> b) -> t CheckM a -> t CheckM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t CheckM SrcLoc
forall a. a -> t CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
patloc
      checkMergeReturn PatBase Info vn (TypeBase shape Diet)
p TypeAliases
_ =
        PatBase Info vn (TypeBase shape Diet)
-> t CheckM (PatBase Info vn (TypeBase shape Diet))
forall a. a -> t CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PatBase Info vn (TypeBase shape Diet)
p

  (param'', (param_cons, _)) <-
    StateT (Aliases, Aliases) CheckM (Pat (TypeBase Size Diet))
-> (Aliases, Aliases)
-> CheckM (Pat (TypeBase Size Diet), (Aliases, Aliases))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (Pat (TypeBase Size Diet)
-> TypeAliases
-> StateT (Aliases, Aliases) CheckM (Pat (TypeBase Size Diet))
forall {t :: (* -> *) -> * -> *} {vn} {shape}.
(MonadTrans t, IsName vn,
 MonadState (Aliases, Aliases) (t CheckM)) =>
PatBase Info vn (TypeBase shape Diet)
-> TypeAliases -> t CheckM (PatBase Info vn (TypeBase shape Diet))
checkMergeReturn Pat (TypeBase Size Diet)
param' TypeAliases
body_als) (Aliases
forall a. Monoid a => a
mempty, Aliases
forall a. Monoid a => a
mempty)

  let body_cons' = Names
body_cons Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> (Alias -> VName) -> Aliases -> Names
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> VName
aliasVar Aliases
param_cons
  if body_cons' == body_cons && patternType param'' == patternType param
    then pure param'
    else convergeLoopParam loop_loc param'' body_cons' body_als

checkLoop :: Loc -> Loop -> CheckM (Loop, TypeAliases)
checkLoop :: Loc -> Loop -> CheckM (Loop, TypeAliases)
checkLoop Loc
loop_loc (Pat (TypeBase Size Diet)
param, LoopInitBase Info VName
arg, LoopFormBase Info VName
form, Size
body) = do
  form' <- LoopFormBase Info VName -> CheckM (LoopFormBase Info VName)
forall e. ASTMappable e => e -> CheckM e
checkSubExps LoopFormBase Info VName
form
  -- We pretend that every part of the loop parameter has a consuming
  -- diet, as we need to allow consumption in the body, which we then
  -- use to infer the proper diet of the parameter.
  ((body', body_cons), body_als) <-
    noConsumable
      . bindingParam (updateParamDiet (const True) param)
      . bindingLoopForm form'
      $ do
        ((body', body_als), body_cons) <- contain $ checkExp body
        pure ((body', body_cons), body_als)
  param' <- convergeLoopParam loop_loc param (M.keysSet body_cons) body_als

  let param_t = Pat (TypeBase Size Diet) -> TypeBase Size Diet
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType Pat (TypeBase Size Diet)
param'
  ((arg', arg_als), arg_cons) <- case arg of
    LoopInitImplicit (Info Size
e) ->
      CheckM (LoopInitBase Info VName, TypeAliases)
-> CheckM ((LoopInitBase Info VName, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (LoopInitBase Info VName, TypeAliases)
 -> CheckM ((LoopInitBase Info VName, TypeAliases), Consumed))
-> CheckM (LoopInitBase Info VName, TypeAliases)
-> CheckM ((LoopInitBase Info VName, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ (Size -> LoopInitBase Info VName)
-> (Size, TypeAliases) -> (LoopInitBase Info VName, TypeAliases)
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 (Info Size -> LoopInitBase Info VName
forall (f :: * -> *) vn. f (ExpBase f vn) -> LoopInitBase f vn
LoopInitImplicit (Info Size -> LoopInitBase Info VName)
-> (Size -> Info Size) -> Size -> LoopInitBase Info VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Size -> Info Size
forall a. a -> Info a
Info) ((Size, TypeAliases) -> (LoopInitBase Info VName, TypeAliases))
-> CheckM (Size, TypeAliases)
-> CheckM (LoopInitBase Info VName, TypeAliases)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Size, TypeAliases)]
-> TypeBase Size Diet -> Size -> CheckM (Size, TypeAliases)
checkArg [] TypeBase Size Diet
param_t Size
e
    LoopInitExplicit Size
e ->
      CheckM (LoopInitBase Info VName, TypeAliases)
-> CheckM ((LoopInitBase Info VName, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (LoopInitBase Info VName, TypeAliases)
 -> CheckM ((LoopInitBase Info VName, TypeAliases), Consumed))
-> CheckM (LoopInitBase Info VName, TypeAliases)
-> CheckM ((LoopInitBase Info VName, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ (Size -> LoopInitBase Info VName)
-> (Size, TypeAliases) -> (LoopInitBase Info VName, TypeAliases)
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 Size -> LoopInitBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopInitBase f vn
LoopInitExplicit ((Size, TypeAliases) -> (LoopInitBase Info VName, TypeAliases))
-> CheckM (Size, TypeAliases)
-> CheckM (LoopInitBase Info VName, TypeAliases)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Size, TypeAliases)]
-> TypeBase Size Diet -> Size -> CheckM (Size, TypeAliases)
checkArg [] TypeBase Size Diet
param_t Size
e
  consumed arg_cons

  let checkFree Doc ()
what Size
e = do
        free_bound <- Size -> CheckM (Map VName TypeAliases)
boundFreeInExp Size
e

        let bad = (VName -> Bool) -> Names -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Consumed -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Consumed
arg_cons) (Names -> Bool)
-> ((a, TypeAliases) -> Names) -> (a, TypeAliases) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliases -> Names
boundAliases (Aliases -> Names)
-> ((a, TypeAliases) -> Aliases) -> (a, TypeAliases) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeAliases -> Aliases
aliases (TypeAliases -> Aliases)
-> ((a, TypeAliases) -> TypeAliases) -> (a, TypeAliases) -> Aliases
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, TypeAliases) -> TypeAliases
forall a b. (a, b) -> b
snd
        forM_ (filter bad $ M.toList free_bound) $ \(VName
v, TypeAliases
_) -> do
          v' <- VName -> CheckM (Doc ())
forall a. VName -> CheckM (Doc a)
describeVar VName
v
          addError loop_loc mempty $
            what
              <+> "uses"
              <+> v'
              <> " (or an alias),"
                </> "but this is consumed by the initial loop argument."

  checkFree "Loop body" body

  case form of
    While Size
cond -> Doc () -> Size -> CheckM ()
checkFree Doc ()
"Loop condition" Size
cond
    LoopFormBase Info VName
_ -> () -> CheckM ()
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  v <- VName "internal_loop_result" <$> incCounter
  modify $ \CheckState
s -> CheckState
s {stateNames = M.insert v (NameLoopRes (srclocOf loop_loc)) $ stateNames s}

  let loopt =
        [Pat (TypeBase Size Diet)]
-> RetTypeBase Size Uniqueness -> StructType
funType [Pat (TypeBase Size Diet)
param'] ([VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (ResType -> RetTypeBase Size Uniqueness)
-> ResType -> RetTypeBase Size Uniqueness
forall a b. (a -> b) -> a -> b
$ TypeBase Size Diet -> ResType
paramToRes TypeBase Size Diet
param_t)
          StructType -> Aliases -> TypeAliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Alias -> Aliases
forall a. a -> Set a
S.singleton (VName -> Alias
AliasFree VName
v)
  pure
    ( (param', arg', form', body'),
      applyArg loopt arg_als `combineAliases` body_als
    )

checkFuncall ::
  (Foldable f) =>
  SrcLoc ->
  Maybe (QualName VName) ->
  TypeAliases ->
  f TypeAliases ->
  CheckM TypeAliases
checkFuncall :: forall (f :: * -> *).
Foldable f =>
SrcLoc
-> Maybe (QualName VName)
-> TypeAliases
-> f TypeAliases
-> CheckM TypeAliases
checkFuncall SrcLoc
loc Maybe (QualName VName)
fname TypeAliases
f_als f TypeAliases
arg_als = do
  v <- Name -> Int -> VName
VName Name
"internal_app_result" (Int -> VName) -> CheckM Int -> CheckM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CheckM Int
incCounter
  modify $ \CheckState
s -> CheckState
s {stateNames = M.insert v (NameAppRes fname loc) $ stateNames s}
  pure $ foldl applyArg (second (S.insert (AliasFree v)) f_als) arg_als

checkExp :: Exp -> CheckM (Exp, TypeAliases)
-- First we have the complicated cases.

--
checkExp :: Size -> CheckM (Size, TypeAliases)
checkExp (AppExp (Apply Size
f NonEmpty (Info (Maybe VName), Size)
args SrcLoc
loc) Info AppRes
appres) = do
  (f', f_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
f
  (args', args_als) <- NE.unzip <$> checkArgs (toRes Nonunique f_als) args
  res_als <- checkFuncall loc (fname f) f_als args_als
  pure (AppExp (Apply f' args' loc) appres, res_als)
  where
    fname :: ExpBase f vn -> Maybe (QualName vn)
fname (Var QualName vn
v f StructType
_ SrcLoc
_) = QualName vn -> Maybe (QualName vn)
forall a. a -> Maybe a
Just QualName vn
v
    fname (AppExp (Apply ExpBase f vn
e NonEmpty (f (Maybe VName), ExpBase f vn)
_ SrcLoc
_) f AppRes
_) = ExpBase f vn -> Maybe (QualName vn)
fname ExpBase f vn
e
    fname ExpBase f vn
_ = Maybe (QualName vn)
forall a. Maybe a
Nothing
    checkArg' :: [(Size, TypeAliases)]
-> Diet -> (Info a, Size) -> CheckM ((Info a, Size), TypeAliases)
checkArg' [(Size, TypeAliases)]
prev Diet
d (Info a
p, Size
e) = do
      (e', e_als) <- [(Size, TypeAliases)]
-> TypeBase Size Diet -> Size -> CheckM (Size, TypeAliases)
checkArg [(Size, TypeAliases)]
prev ((NoUniqueness -> Diet) -> StructType -> TypeBase Size Diet
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Diet -> NoUniqueness -> Diet
forall a b. a -> b -> a
const Diet
d) (Size -> StructType
typeOf Size
e)) Size
e
      pure ((Info p, e'), e_als)

    checkArgs :: TypeBase dim Uniqueness
-> NonEmpty (Info a, Size)
-> CheckM (NonEmpty ((Info a, Size), TypeAliases))
checkArgs (Scalar (Arrow Uniqueness
_ PName
_ Diet
d TypeBase dim NoUniqueness
_ (RetType [VName]
_ TypeBase dim Uniqueness
rt))) ((Info a, Size)
x NE.:| [(Info a, Size)]
args') = do
      -- Note Futhark uses right-to-left evaluation of applications.
      args'' <- CheckM [((Info a, Size), TypeAliases)]
-> (NonEmpty (Info a, Size)
    -> CheckM [((Info a, Size), TypeAliases)])
-> Maybe (NonEmpty (Info a, Size))
-> CheckM [((Info a, Size), TypeAliases)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([((Info a, Size), TypeAliases)]
-> CheckM [((Info a, Size), TypeAliases)]
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []) ((NonEmpty ((Info a, Size), TypeAliases)
 -> [((Info a, Size), TypeAliases)])
-> CheckM (NonEmpty ((Info a, Size), TypeAliases))
-> CheckM [((Info a, Size), TypeAliases)]
forall a b. (a -> b) -> CheckM a -> CheckM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty ((Info a, Size), TypeAliases)
-> [((Info a, Size), TypeAliases)]
forall a. NonEmpty a -> [a]
NE.toList (CheckM (NonEmpty ((Info a, Size), TypeAliases))
 -> CheckM [((Info a, Size), TypeAliases)])
-> (NonEmpty (Info a, Size)
    -> CheckM (NonEmpty ((Info a, Size), TypeAliases)))
-> NonEmpty (Info a, Size)
-> CheckM [((Info a, Size), TypeAliases)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase dim Uniqueness
-> NonEmpty (Info a, Size)
-> CheckM (NonEmpty ((Info a, Size), TypeAliases))
checkArgs TypeBase dim Uniqueness
rt) (Maybe (NonEmpty (Info a, Size))
 -> CheckM [((Info a, Size), TypeAliases)])
-> Maybe (NonEmpty (Info a, Size))
-> CheckM [((Info a, Size), TypeAliases)]
forall a b. (a -> b) -> a -> b
$ [(Info a, Size)] -> Maybe (NonEmpty (Info a, Size))
forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty [(Info a, Size)]
args'
      (x', x_als) <- checkArg' (map (first snd) args'') d x
      pure $ (x', x_als) NE.:| args''
    checkArgs TypeBase dim Uniqueness
t NonEmpty (Info a, Size)
_ =
      String -> CheckM (NonEmpty ((Info a, Size), TypeAliases))
forall a. HasCallStack => String -> a
error (String -> CheckM (NonEmpty ((Info a, Size), TypeAliases)))
-> String -> CheckM (NonEmpty ((Info a, Size), TypeAliases))
forall a b. (a -> b) -> a -> b
$ String
"checkArgs: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> TypeBase dim Uniqueness -> String
forall a. Pretty a => a -> String
prettyString TypeBase dim Uniqueness
t

--
checkExp (AppExp (Loop [VName]
sparams Pat (TypeBase Size Diet)
pat LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Size
body SrcLoc
loc) Info AppRes
appres) = do
  ((pat', loopinit', form', body'), als) <-
    Loc -> Loop -> CheckM (Loop, TypeAliases)
checkLoop (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Pat (TypeBase Size Diet)
pat, LoopInitBase Info VName
loopinit, LoopFormBase Info VName
form, Size
body)
  pure
    ( AppExp (Loop sparams pat' loopinit' form' body' loc) appres,
      als
    )

--
checkExp (AppExp (LetPat [SizeBinder VName]
sizes Pat StructType
p Size
e Size
body SrcLoc
loc) Info AppRes
appres) = do
  ((e', e_als), e_cons) <- CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (Size, TypeAliases)
 -> CheckM ((Size, TypeAliases), Consumed))
-> CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  consumed e_cons
  let e_t = Size -> StructType
typeOf Size
e'
  when (e_cons /= mempty && not (orderZero e_t)) $
    addError (locOf e) mempty $
      "Let-bound expression of higher-order type"
        </> indent 2 (pretty e_t)
        </> "contains consumption, which is not allowed."
  bindingPat p e_als $ do
    (body', body_als) <- checkExp body
    pure
      ( AppExp (LetPat sizes p e' body' loc) appres,
        body_als
      )

--
checkExp (AppExp (If Size
cond Size
te Size
fe SrcLoc
loc) Info AppRes
appres) = do
  (cond', _) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
cond
  ((te', te_als), te_cons) <- contain $ checkExp te
  ((fe', fe_als), fe_cons) <- contain $ checkExp fe
  let all_cons = Consumed
te_cons Consumed -> Consumed -> Consumed
forall a. Semigroup a => a -> a -> a
<> Consumed
fe_cons
      notConsumed = Bool -> Bool
not (Bool -> Bool) -> (Alias -> Bool) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Consumed -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Consumed
all_cons) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar
      comb_als = (Aliases -> Aliases) -> TypeAliases -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Alias -> Bool) -> Aliases -> Aliases
forall a. (a -> Bool) -> Set a -> Set a
S.filter Alias -> Bool
notConsumed) (TypeAliases -> TypeAliases) -> TypeAliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ TypeAliases
te_als TypeAliases -> TypeAliases -> TypeAliases
`combineAliases` TypeAliases
fe_als
  consumed all_cons
  pure
    ( AppExp (If cond' te' fe' loc) appres,
      appResType (unInfo appres) `setAliases` mempty `combineAliases` comb_als
    )

--
checkExp (AppExp (Match Size
cond NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
appres) = do
  (cond', cond_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
cond
  ((cs', cs_als), cs_cons) <-
    first NE.unzip . NE.unzip <$> mapM (checkCase cond_als) cs
  let all_cons = NonEmpty Consumed -> Consumed
forall m. Monoid m => NonEmpty m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold NonEmpty Consumed
cs_cons
      notConsumed = Bool -> Bool
not (Bool -> Bool) -> (Alias -> Bool) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Consumed -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Consumed
all_cons) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar
      comb_als = (Aliases -> Aliases) -> TypeAliases -> TypeAliases
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Alias -> Bool) -> Aliases -> Aliases
forall a. (a -> Bool) -> Set a -> Set a
S.filter Alias -> Bool
notConsumed) (TypeAliases -> TypeAliases) -> TypeAliases -> TypeAliases
forall a b. (a -> b) -> a -> b
$ (TypeAliases -> TypeAliases -> TypeAliases)
-> NonEmpty TypeAliases -> TypeAliases
forall a. (a -> a -> a) -> NonEmpty a -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 TypeAliases -> TypeAliases -> TypeAliases
combineAliases NonEmpty TypeAliases
cs_als
  consumed all_cons
  pure
    ( AppExp (Match cond' cs' loc) appres,
      appResType (unInfo appres) `setAliases` mempty `combineAliases` comb_als
    )
  where
    checkCase :: TypeAliases
-> CaseBase Info VName
-> CheckM ((CaseBase Info VName, TypeAliases), Consumed)
checkCase TypeAliases
cond_als (CasePat Pat StructType
p Size
body SrcLoc
caseloc) =
      CheckM (CaseBase Info VName, TypeAliases)
-> CheckM ((CaseBase Info VName, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (CaseBase Info VName, TypeAliases)
 -> CheckM ((CaseBase Info VName, TypeAliases), Consumed))
-> CheckM (CaseBase Info VName, TypeAliases)
-> CheckM ((CaseBase Info VName, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ Pat StructType
-> TypeAliases
-> CheckM (CaseBase Info VName, TypeAliases)
-> CheckM (CaseBase Info VName, TypeAliases)
forall a.
Pat StructType
-> TypeAliases
-> CheckM (a, TypeAliases)
-> CheckM (a, TypeAliases)
bindingPat Pat StructType
p TypeAliases
cond_als (CheckM (CaseBase Info VName, TypeAliases)
 -> CheckM (CaseBase Info VName, TypeAliases))
-> CheckM (CaseBase Info VName, TypeAliases)
-> CheckM (CaseBase Info VName, TypeAliases)
forall a b. (a -> b) -> a -> b
$ do
        (body', body_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
body
        pure (CasePat p body' caseloc, body_als)

--
checkExp (AppExp (LetFun VName
fname ([TypeParamBase VName]
typarams, [Pat (TypeBase Size Diet)]
params, Maybe (TypeExp Size VName)
te, Info (RetType [VName]
ext ResType
ret), Size
funbody) Size
letbody SrcLoc
loc) Info AppRes
appres) = do
  ((ret', funbody'), ftype) <- [Pat (TypeBase Size Diet)]
-> CheckM ((ResType, Size), TypeAliases)
-> CheckM ((ResType, Size), TypeAliases)
forall a.
[Pat (TypeBase Size Diet)]
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParams [Pat (TypeBase Size Diet)]
params (CheckM ((ResType, Size), TypeAliases)
 -> CheckM ((ResType, Size), TypeAliases))
-> CheckM ((ResType, Size), TypeAliases)
-> CheckM ((ResType, Size), TypeAliases)
forall a b. (a -> b) -> a -> b
$ do
    -- Throw away the consumption - it can refer only to the parameters
    -- anyway.
    ((funbody', funbody_als), _body_cons) <- CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (Size, TypeAliases)
 -> CheckM ((Size, TypeAliases), Consumed))
-> CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ Size -> CheckM (Size, TypeAliases)
checkExp Size
funbody
    checkReturnAlias loc params ret funbody_als
    checkGlobalAliases loc params funbody_als
    free_bound <- boundFreeInExp funbody
    let ret' = [Pat (TypeBase Size Diet)] -> ResType -> TypeAliases -> ResType
inferReturnUniqueness [Pat (TypeBase Size Diet)]
params ResType
ret TypeAliases
funbody_als
        als = (TypeAliases -> Aliases) -> [TypeAliases] -> Aliases
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeAliases -> Aliases
aliases (Map VName TypeAliases -> [TypeAliases]
forall k a. Map k a -> [a]
M.elems Map VName TypeAliases
free_bound)
        ftype = [Pat (TypeBase Size Diet)]
-> RetTypeBase Size Uniqueness -> StructType
funType [Pat (TypeBase Size Diet)]
params ([VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext ResType
ret') StructType -> Aliases -> TypeAliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliases
als
    pure ((ret', funbody'), ftype)
  (letbody', letbody_als) <- bindingFun fname ftype $ checkExp letbody
  pure
    ( AppExp (LetFun fname (typarams, params, te, Info (RetType ext ret'), funbody') letbody' loc) appres,
      letbody_als
    )

--
checkExp (AppExp (BinOp (QualName VName
op, SrcLoc
oploc) Info StructType
opt (Size
x, Info (Maybe VName)
xp) (Size
y, Info (Maybe VName)
yp) SrcLoc
loc) Info AppRes
appres) = do
  op_als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
oploc) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) (Info StructType -> StructType
forall a. Info a -> a
unInfo Info StructType
opt)
  let at1 : at2 : _ = fst $ unfoldFunType op_als
  (x', x_als) <- checkArg [] at1 x
  (y', y_als) <- checkArg [(x', x_als)] at2 y
  res_als <- checkFuncall loc (Just op) op_als [x_als, y_als]
  pure
    ( AppExp (BinOp (op, oploc) opt (x', xp) (y', yp) loc) appres,
      res_als
    )

--
checkExp e :: Size
e@(Lambda [Pat (TypeBase Size Diet)]
params Size
body Maybe (TypeExp Size VName)
te (Info (RetType [VName]
ext ResType
ret)) SrcLoc
loc) =
  [Pat (TypeBase Size Diet)]
-> CheckM (Size, TypeAliases) -> CheckM (Size, TypeAliases)
forall a.
[Pat (TypeBase Size Diet)]
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParams [Pat (TypeBase Size Diet)]
params (CheckM (Size, TypeAliases) -> CheckM (Size, TypeAliases))
-> CheckM (Size, TypeAliases) -> CheckM (Size, TypeAliases)
forall a b. (a -> b) -> a -> b
$ do
    -- Throw away the consumption - it can refer only to the parameters
    -- anyway.
    ((body', body_als), _body_cons) <- CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a. CheckM a -> CheckM (a, Consumed)
contain (CheckM (Size, TypeAliases)
 -> CheckM ((Size, TypeAliases), Consumed))
-> CheckM (Size, TypeAliases)
-> CheckM ((Size, TypeAliases), Consumed)
forall a b. (a -> b) -> a -> b
$ Size -> CheckM (Size, TypeAliases)
checkExp Size
body
    checkReturnAlias loc params ret body_als
    checkGlobalAliases loc params body_als
    free_bound <- boundFreeInExp e
    let ret' = [Pat (TypeBase Size Diet)] -> ResType -> TypeAliases -> ResType
inferReturnUniqueness [Pat (TypeBase Size Diet)]
params ResType
ret TypeAliases
body_als
        als = (TypeAliases -> Aliases) -> [TypeAliases] -> Aliases
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeAliases -> Aliases
aliases (Map VName TypeAliases -> [TypeAliases]
forall k a. Map k a -> [a]
M.elems Map VName TypeAliases
free_bound)
        ftype = [Pat (TypeBase Size Diet)]
-> RetTypeBase Size Uniqueness -> StructType
funType [Pat (TypeBase Size Diet)]
params ([VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext ResType
ret') StructType -> Aliases -> TypeAliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliases
als
    pure
      ( Lambda params body' te (Info (RetType ext ret')) loc,
        ftype
      )

--
checkExp (AppExp (LetWith IdentBase Info VName StructType
dst IdentBase Info VName StructType
src SliceBase Info VName
slice Size
ve Size
body SrcLoc
loc) Info AppRes
appres) = do
  src_als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (IdentBase Info VName StructType -> Loc
forall a. Located a => a -> Loc
locOf IdentBase Info VName StructType
dst) (IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase Info VName StructType
src) (Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> StructType) -> Info StructType -> StructType
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> Info StructType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType IdentBase Info VName StructType
src)
  slice' <- checkSubExps slice
  (ve', ve_als) <- checkExp ve
  consume (locOf src) (identName src) (unInfo (identType src))
  overlapCheck (locOf ve) (src, src_als) (ve', ve_als)
  (body', body_als) <- bindingIdent Consume dst $ checkExp body
  pure (AppExp (LetWith dst src slice' ve' body' loc) appres, body_als)

--
checkExp (Update Size
src SliceBase Info VName
slice Size
ve SrcLoc
loc) = do
  slice' <- SliceBase Info VName -> CheckM (SliceBase Info VName)
forall e. ASTMappable e => e -> CheckM e
checkSubExps SliceBase Info VName
slice
  (ve', ve_als) <- checkExp ve
  (src', src_als) <- checkExp src
  overlapCheck (locOf ve) (src', src_als) (ve', ve_als)
  consumeAliases (locOf loc) $ aliases src_als
  pure (Update src' slice' ve' loc, second (const mempty) src_als)

-- Cases that simply propagate aliases directly.
checkExp (Var QualName VName
v (Info StructType
t) SrcLoc
loc) = do
  als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) StructType
t
  checkIfConsumed (locOf loc) (aliases als)
  pure (Var v (Info t) loc, als)
checkExp (OpSection QualName VName
v (Info StructType
t) SrcLoc
loc) = do
  als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) StructType
t
  checkIfConsumed (locOf loc) (aliases als)
  pure (OpSection v (Info t) loc, als)
checkExp (OpSectionLeft QualName VName
op Info StructType
ftype Size
arg (Info (PName, TypeBase Size Diet, Maybe VName),
 Info (PName, TypeBase Size Diet))
arginfo (Info (RetTypeBase Size Uniqueness), Info [VName])
retinfo SrcLoc
loc) = do
  let (Info (PName, TypeBase Size Diet, Maybe VName)
_, Info (PName
pn, TypeBase Size Diet
pt2)) = (Info (PName, TypeBase Size Diet, Maybe VName),
 Info (PName, TypeBase Size Diet))
arginfo
      (Info RetTypeBase Size Uniqueness
ret, Info [VName]
_) = (Info (RetTypeBase Size Uniqueness), Info [VName])
retinfo
  als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) (Info StructType -> StructType
forall a. Info a -> a
unInfo Info StructType
ftype)
  (arg', arg_als) <- checkExp arg
  pure
    ( OpSectionLeft op ftype arg' arginfo retinfo loc,
      Scalar $ Arrow (aliases arg_als <> aliases als) pn (diet pt2) (toStruct pt2) ret
    )
checkExp (OpSectionRight QualName VName
op Info StructType
ftype Size
arg (Info (PName, TypeBase Size Diet),
 Info (PName, TypeBase Size Diet, Maybe VName))
arginfo Info (RetTypeBase Size Uniqueness)
retinfo SrcLoc
loc) = do
  let (Info (PName
pn, TypeBase Size Diet
pt2), Info (PName, TypeBase Size Diet, Maybe VName)
_) = (Info (PName, TypeBase Size Diet),
 Info (PName, TypeBase Size Diet, Maybe VName))
arginfo
      Info RetTypeBase Size Uniqueness
ret = Info (RetTypeBase Size Uniqueness)
retinfo
  als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) (Info StructType -> StructType
forall a. Info a -> a
unInfo Info StructType
ftype)
  (arg', arg_als) <- checkExp arg
  pure
    ( OpSectionRight op ftype arg' arginfo retinfo loc,
      Scalar $ Arrow (aliases arg_als <> aliases als) pn (diet pt2) (toStruct pt2) ret
    )
checkExp (IndexSection SliceBase Info VName
slice Info StructType
t SrcLoc
loc) = do
  slice' <- SliceBase Info VName -> CheckM (SliceBase Info VName)
forall e. ASTMappable e => e -> CheckM e
checkSubExps SliceBase Info VName
slice
  pure (IndexSection slice' t loc, unInfo t `setAliases` mempty)
checkExp (ProjectSection [Name]
fs Info StructType
t SrcLoc
loc) = do
  (Size, TypeAliases) -> CheckM (Size, TypeAliases)
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Name] -> Info StructType -> SrcLoc -> Size
forall (f :: * -> *) vn.
[Name] -> f StructType -> SrcLoc -> ExpBase f vn
ProjectSection [Name]
fs Info StructType
t SrcLoc
loc, Info StructType -> StructType
forall a. Info a -> a
unInfo Info StructType
t StructType -> Aliases -> TypeAliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliases
forall a. Monoid a => a
mempty)
checkExp (Coerce Size
e TypeExp Size VName
te Info StructType
t SrcLoc
loc) = do
  (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  pure (Coerce e' te t loc, e_als)
checkExp (Ascript Size
e TypeExp Size VName
te SrcLoc
loc) = do
  (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  pure (Ascript e' te loc, e_als)
checkExp (AppExp (Index Size
v SliceBase Info VName
slice SrcLoc
loc) Info AppRes
appres) = do
  (v', v_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
v
  slice' <- checkSubExps slice
  pure
    ( AppExp (Index v' slice' loc) appres,
      appResType (unInfo appres) `setAliases` aliases v_als
    )
checkExp (Assert Size
e1 Size
e2 Info Text
t SrcLoc
loc) = do
  (e1', _) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e1
  (e2', e2_als) <- checkExp e2
  pure (Assert e1' e2' t loc, e2_als)
checkExp (Parens Size
e SrcLoc
loc) = do
  (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  pure (Parens e' loc, e_als)
checkExp (QualParens (QualName VName, SrcLoc)
v Size
e SrcLoc
loc) = do
  (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  pure (QualParens v e' loc, e_als)
checkExp (Attr AttrInfo VName
attr Size
e SrcLoc
loc) = do
  (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  pure (Attr attr e' loc, e_als)
checkExp (Project Name
name Size
e Info StructType
t SrcLoc
loc) = do
  (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
  pure
    ( Project name e' t loc,
      case e_als of
        Scalar (Record Map Name TypeAliases
fs)
          | Just TypeAliases
name_als <- Name -> Map Name TypeAliases -> Maybe TypeAliases
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
name Map Name TypeAliases
fs -> TypeAliases
name_als
        TypeAliases
_ -> String -> TypeAliases
forall a. HasCallStack => String -> a
error (String -> TypeAliases) -> String -> TypeAliases
forall a b. (a -> b) -> a -> b
$ String
"checkExp Project: bad type " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> TypeAliases -> String
forall a. Pretty a => a -> String
prettyString TypeAliases
e_als
    )
checkExp (TupLit [Size]
es SrcLoc
loc) = do
  (es', es_als) <- (Size -> CheckM (Size, TypeAliases))
-> [Size] -> CheckM ([Size], [TypeAliases])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Size -> CheckM (Size, TypeAliases)
checkExp [Size]
es
  pure (TupLit es' loc, Scalar $ tupleRecord es_als)
checkExp (Constr Name
name [Size]
es Info StructType
t SrcLoc
loc) = do
  (es', es_als) <- (Size -> CheckM (Size, TypeAliases))
-> [Size] -> CheckM ([Size], [TypeAliases])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Size -> CheckM (Size, TypeAliases)
checkExp [Size]
es
  pure
    ( Constr name es' t loc,
      case unInfo t of
        Scalar (Sum Map Name [StructType]
cs) ->
          ScalarTypeBase Size Aliases -> TypeAliases
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Size Aliases -> TypeAliases)
-> (Map Name [TypeAliases] -> ScalarTypeBase Size Aliases)
-> Map Name [TypeAliases]
-> TypeAliases
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeAliases] -> ScalarTypeBase Size Aliases
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeAliases] -> ScalarTypeBase Size Aliases)
-> (Map Name [TypeAliases] -> Map Name [TypeAliases])
-> Map Name [TypeAliases]
-> ScalarTypeBase Size Aliases
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name
-> [TypeAliases]
-> Map Name [TypeAliases]
-> Map Name [TypeAliases]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
name [TypeAliases]
es_als (Map Name [TypeAliases] -> TypeAliases)
-> Map Name [TypeAliases] -> TypeAliases
forall a b. (a -> b) -> a -> b
$
            ([StructType] -> [TypeAliases])
-> Map Name [StructType] -> Map Name [TypeAliases]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> TypeAliases) -> [StructType] -> [TypeAliases]
forall a b. (a -> b) -> [a] -> [b]
map (StructType -> Aliases -> TypeAliases
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliases
forall a. Monoid a => a
mempty)) Map Name [StructType]
cs
        StructType
t' -> String -> TypeAliases
forall a. HasCallStack => String -> a
error (String -> TypeAliases) -> String -> TypeAliases
forall a b. (a -> b) -> a -> b
$ String
"checkExp Constr: bad type " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> StructType -> String
forall a. Pretty a => a -> String
prettyString StructType
t'
    )
checkExp (RecordUpdate Size
src [Name]
fields Size
ve Info StructType
t SrcLoc
loc) = do
  (src', src_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
src
  (ve', ve_als) <- checkExp ve
  pure
    ( RecordUpdate src' fields ve' t loc,
      setFieldAliases ve_als fields src_als
    )
checkExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
  (fs', fs_als) <- (FieldBase Info VName
 -> CheckM (FieldBase Info VName, (Name, TypeAliases)))
-> [FieldBase Info VName]
-> CheckM ([FieldBase Info VName], [(Name, TypeAliases)])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM FieldBase Info VName
-> CheckM (FieldBase Info VName, (Name, TypeAliases))
checkField [FieldBase Info VName]
fs
  pure (RecordLit fs' loc, Scalar $ Record $ M.fromList fs_als)
  where
    checkField :: FieldBase Info VName
-> CheckM (FieldBase Info VName, (Name, TypeAliases))
checkField (RecordFieldExplicit L Name
name Size
e SrcLoc
floc) = do
      (e', e_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
e
      pure (RecordFieldExplicit name e' floc, (unLoc name, e_als))
    checkField (RecordFieldImplicit L VName
name Info StructType
t SrcLoc
floc) = do
      name_als <- Loc -> VName -> StructType -> CheckM TypeAliases
observeVar (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
floc) (L VName -> VName
forall a. L a -> a
unLoc L VName
name) (StructType -> CheckM TypeAliases)
-> StructType -> CheckM TypeAliases
forall a b. (a -> b) -> a -> b
$ Info StructType -> StructType
forall a. Info a -> a
unInfo Info StructType
t
      pure (RecordFieldImplicit name t floc, (baseName (unLoc name), name_als))

-- Cases that create alias-free values.
checkExp e :: Size
e@(AppExp Range {} Info AppRes
_) = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@IntLit {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@FloatLit {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@Literal {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@StringLit {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@ArrayVal {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@ArrayLit {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@Negate {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@Not {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e
checkExp e :: Size
e@Hole {} = Size -> CheckM (Size, TypeAliases)
noAliases Size
e

checkGlobalAliases :: SrcLoc -> [Pat ParamType] -> TypeAliases -> CheckM ()
checkGlobalAliases :: SrcLoc -> [Pat (TypeBase Size Diet)] -> TypeAliases -> CheckM ()
checkGlobalAliases SrcLoc
loc [Pat (TypeBase Size Diet)]
params TypeAliases
body_t = do
  vtable <- (CheckEnv -> Map VName (Entry TypeAliases))
-> CheckM (Map VName (Entry TypeAliases))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CheckEnv -> Map VName (Entry TypeAliases)
envVtable
  let global = (VName -> Map VName (Entry TypeAliases) -> Bool)
-> Map VName (Entry TypeAliases) -> VName -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> Map VName (Entry TypeAliases) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Map VName (Entry TypeAliases)
vtable
  unless (null params) $ forM_ (boundAliases $ arrayAliases body_t) $ \VName
v ->
    Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (VName -> Bool
global VName
v) (CheckM () -> CheckM ())
-> (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc () -> CheckM ()) -> (Doc () -> Doc ()) -> Doc () -> CheckM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
withIndexLink Doc ()
"alias-free-variable" (Doc () -> CheckM ()) -> Doc () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
      Doc ()
"Function result aliases the free variable "
        Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes (VName -> Doc ()
forall a. VName -> Doc a
forall v a. IsName v => v -> Doc a
prettyName VName
v)
        Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
"."
          Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc ()
"Use"
          Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
dquotes Doc ()
"copy"
          Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
"to break the aliasing."

-- | Type-check a value definition.  This also infers a new return
-- type that may be more unique than previously.
checkValDef ::
  (VName, [Pat ParamType], Exp, ResRetType, Maybe (TypeExp Exp VName), SrcLoc) ->
  ((Exp, ResRetType), [TypeError])
checkValDef :: (VName, [Pat (TypeBase Size Diet)], Size,
 RetTypeBase Size Uniqueness, Maybe (TypeExp Size VName), SrcLoc)
-> ((Size, RetTypeBase Size Uniqueness), [TypeError])
checkValDef (VName
_fname, [Pat (TypeBase Size Diet)]
params, Size
body, RetType [VName]
ext ResType
ret, Maybe (TypeExp Size VName)
retdecl, SrcLoc
loc) = Loc
-> CheckM (Size, RetTypeBase Size Uniqueness)
-> ((Size, RetTypeBase Size Uniqueness), [TypeError])
forall a. Loc -> CheckM a -> (a, [TypeError])
runCheckM (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (CheckM (Size, RetTypeBase Size Uniqueness)
 -> ((Size, RetTypeBase Size Uniqueness), [TypeError]))
-> CheckM (Size, RetTypeBase Size Uniqueness)
-> ((Size, RetTypeBase Size Uniqueness), [TypeError])
forall a b. (a -> b) -> a -> b
$ do
  (((Size, RetTypeBase Size Uniqueness), TypeAliases)
 -> (Size, RetTypeBase Size Uniqueness))
-> CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
-> CheckM (Size, RetTypeBase Size Uniqueness)
forall a b. (a -> b) -> CheckM a -> CheckM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Size, RetTypeBase Size Uniqueness), TypeAliases)
-> (Size, RetTypeBase Size Uniqueness)
forall a b. (a, b) -> a
fst (CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
 -> CheckM (Size, RetTypeBase Size Uniqueness))
-> (CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
    -> CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases))
-> CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
-> CheckM (Size, RetTypeBase Size Uniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Pat (TypeBase Size Diet)]
-> CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
-> CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
forall a.
[Pat (TypeBase Size Diet)]
-> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
bindingParams [Pat (TypeBase Size Diet)]
params (CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
 -> CheckM (Size, RetTypeBase Size Uniqueness))
-> CheckM ((Size, RetTypeBase Size Uniqueness), TypeAliases)
-> CheckM (Size, RetTypeBase Size Uniqueness)
forall a b. (a -> b) -> a -> b
$ do
    (body', body_als) <- Size -> CheckM (Size, TypeAliases)
checkExp Size
body
    checkReturnAlias loc params ret body_als
    checkGlobalAliases loc params body_als
    -- If the user did not provide an annotation (meaning the return
    -- type is fully inferred), we infer the uniqueness.  Otherwise,
    -- we go with whatever they wanted.  This lets the user define
    -- non-unique return types even if the body actually has no
    -- aliases.
    ret' <- case retdecl of
      Just TypeExp Size VName
retdecl' -> do
        Bool -> CheckM () -> CheckM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([Pat (TypeBase Size Diet)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Pat (TypeBase Size Diet)]
params Bool -> Bool -> Bool
&& ResType -> Bool
forall shape. TypeBase shape Uniqueness -> Bool
unique ResType
ret) (CheckM () -> CheckM ()) -> CheckM () -> CheckM ()
forall a b. (a -> b) -> a -> b
$
          TypeExp Size VName -> Notes -> Doc () -> CheckM ()
forall loc. Located loc => loc -> Notes -> Doc () -> CheckM ()
addError TypeExp Size VName
retdecl' Notes
forall a. Monoid a => a
mempty Doc ()
"A top-level constant cannot have a unique type."
        RetTypeBase Size Uniqueness -> CheckM (RetTypeBase Size Uniqueness)
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase Size Uniqueness
 -> CheckM (RetTypeBase Size Uniqueness))
-> RetTypeBase Size Uniqueness
-> CheckM (RetTypeBase Size Uniqueness)
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext ResType
ret
      Maybe (TypeExp Size VName)
Nothing ->
        RetTypeBase Size Uniqueness -> CheckM (RetTypeBase Size Uniqueness)
forall a. a -> CheckM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase Size Uniqueness
 -> CheckM (RetTypeBase Size Uniqueness))
-> RetTypeBase Size Uniqueness
-> CheckM (RetTypeBase Size Uniqueness)
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> RetTypeBase Size Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext (ResType -> RetTypeBase Size Uniqueness)
-> ResType -> RetTypeBase Size Uniqueness
forall a b. (a -> b) -> a -> b
$ [Pat (TypeBase Size Diet)] -> ResType -> TypeAliases -> ResType
inferReturnUniqueness [Pat (TypeBase Size Diet)]
params ResType
ret TypeAliases
body_als
    pure
      ( (body', ret'),
        body_als -- Don't matter.
      )
{-# NOINLINE checkValDef #-}