{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

-- | Type checker building blocks that do not involve unification.
module Language.Futhark.TypeChecker.Types
  ( checkTypeExp,
    renameRetType,
    unifyTypesU,
    subtypeOf,
    subuniqueOf,
    checkForDuplicateNames,
    checkTypeParams,
    typeParamToArg,
    Subst (..),
    substFromAbbr,
    TypeSubs,
    Substitutable (..),
    substTypesAny,
  )
where

import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.List (foldl', sort, unzip4, (\\))
import qualified Data.Map.Strict as M
import Data.Maybe
import Futhark.Util (nubOrd)
import Futhark.Util.Pretty hiding ((<|>))
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Monad

-- | @unifyTypes uf t1 t2@ attempts to unify @t1@ and @t2@.  If
-- unification cannot happen, 'Nothing' is returned, otherwise a type
-- that combines the aliasing of @t1@ and @t2@ is returned.
-- Uniqueness is unified with @uf@.  Assumes sizes already match, and
-- always picks the size of the leftmost type.
unifyTypesU ::
  (Monoid als, ArrayDim dim) =>
  (Uniqueness -> Uniqueness -> Maybe Uniqueness) ->
  TypeBase dim als ->
  TypeBase dim als ->
  Maybe (TypeBase dim als)
unifyTypesU :: (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
uf (Array als
als1 Uniqueness
u1 ScalarTypeBase dim ()
et1 ShapeDecl dim
shape1) (Array als
als2 Uniqueness
u2 ScalarTypeBase dim ()
et2 ShapeDecl dim
_shape2) =
  als
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim als
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array (als
als1 als -> als -> als
forall a. Semigroup a => a -> a -> a
<> als
als2) (Uniqueness
 -> ScalarTypeBase dim () -> ShapeDecl dim -> TypeBase dim als)
-> Maybe Uniqueness
-> Maybe
     (ScalarTypeBase dim () -> ShapeDecl dim -> TypeBase dim als)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Uniqueness -> Uniqueness -> Maybe Uniqueness
uf Uniqueness
u1 Uniqueness
u2
    Maybe (ScalarTypeBase dim () -> ShapeDecl dim -> TypeBase dim als)
-> Maybe (ScalarTypeBase dim ())
-> Maybe (ShapeDecl dim -> TypeBase dim als)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> ScalarTypeBase dim ()
-> ScalarTypeBase dim ()
-> Maybe (ScalarTypeBase dim ())
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> ScalarTypeBase dim als
-> ScalarTypeBase dim als
-> Maybe (ScalarTypeBase dim als)
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
uf ScalarTypeBase dim ()
et1 ScalarTypeBase dim ()
et2
    Maybe (ShapeDecl dim -> TypeBase dim als)
-> Maybe (ShapeDecl dim) -> Maybe (TypeBase dim als)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ShapeDecl dim -> Maybe (ShapeDecl dim)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ShapeDecl dim
shape1
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
uf (Scalar ScalarTypeBase dim als
t1) (Scalar ScalarTypeBase dim als
t2) = ScalarTypeBase dim als -> TypeBase dim als
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim als -> TypeBase dim als)
-> Maybe (ScalarTypeBase dim als) -> Maybe (TypeBase dim als)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> ScalarTypeBase dim als
-> ScalarTypeBase dim als
-> Maybe (ScalarTypeBase dim als)
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> ScalarTypeBase dim als
-> ScalarTypeBase dim als
-> Maybe (ScalarTypeBase dim als)
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
uf ScalarTypeBase dim als
t1 ScalarTypeBase dim als
t2
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
_ TypeBase dim als
_ TypeBase dim als
_ = Maybe (TypeBase dim als)
forall a. Maybe a
Nothing

unifyScalarTypes ::
  (Monoid als, ArrayDim dim) =>
  (Uniqueness -> Uniqueness -> Maybe Uniqueness) ->
  ScalarTypeBase dim als ->
  ScalarTypeBase dim als ->
  Maybe (ScalarTypeBase dim als)
unifyScalarTypes :: (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> ScalarTypeBase dim als
-> ScalarTypeBase dim als
-> Maybe (ScalarTypeBase dim als)
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
_ (Prim PrimType
t1) (Prim PrimType
t2)
  | PrimType
t1 PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
t2 = ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als)
forall a. a -> Maybe a
Just (ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als))
-> ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als)
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim als
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t1
  | Bool
otherwise = Maybe (ScalarTypeBase dim als)
forall a. Maybe a
Nothing
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
uf (TypeVar als
als1 Uniqueness
u1 TypeName
tv1 [TypeArg dim]
targs1) (TypeVar als
als2 Uniqueness
u2 TypeName
tv2 [TypeArg dim]
targs2)
  | TypeName
tv1 TypeName -> TypeName -> Bool
forall a. Eq a => a -> a -> Bool
== TypeName
tv2 = do
    Uniqueness
u3 <- Uniqueness -> Uniqueness -> Maybe Uniqueness
uf Uniqueness
u1 Uniqueness
u2
    [TypeArg dim]
targs3 <- (TypeArg dim -> TypeArg dim -> Maybe (TypeArg dim))
-> [TypeArg dim] -> [TypeArg dim] -> Maybe [TypeArg dim]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeArg dim -> TypeArg dim -> Maybe (TypeArg dim)
forall dim.
ArrayDim dim =>
TypeArg dim -> TypeArg dim -> Maybe (TypeArg dim)
unifyTypeArgs [TypeArg dim]
targs1 [TypeArg dim]
targs2
    ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als)
forall a. a -> Maybe a
Just (ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als))
-> ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als)
forall a b. (a -> b) -> a -> b
$ als
-> Uniqueness
-> TypeName
-> [TypeArg dim]
-> ScalarTypeBase dim als
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar (als
als1 als -> als -> als
forall a. Semigroup a => a -> a -> a
<> als
als2) Uniqueness
u3 TypeName
tv1 [TypeArg dim]
targs3
  | Bool
otherwise = Maybe (ScalarTypeBase dim als)
forall a. Maybe a
Nothing
  where
    unifyTypeArgs :: TypeArg dim -> TypeArg dim -> Maybe (TypeArg dim)
unifyTypeArgs (TypeArgDim dim
d1 SrcLoc
loc) (TypeArgDim dim
_d2 SrcLoc
_) =
      TypeArg dim -> Maybe (TypeArg dim)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeArg dim -> Maybe (TypeArg dim))
-> TypeArg dim -> Maybe (TypeArg dim)
forall a b. (a -> b) -> a -> b
$ dim -> SrcLoc -> TypeArg dim
forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim dim
d1 SrcLoc
loc
    unifyTypeArgs (TypeArgType TypeBase dim ()
t1 SrcLoc
loc) (TypeArgType TypeBase dim ()
t2 SrcLoc
_) =
      TypeBase dim () -> SrcLoc -> TypeArg dim
forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (TypeBase dim () -> SrcLoc -> TypeArg dim)
-> Maybe (TypeBase dim ()) -> Maybe (SrcLoc -> TypeArg dim)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim () -> TypeBase dim () -> Maybe (TypeBase dim ())
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
uf TypeBase dim ()
t1 TypeBase dim ()
t2 Maybe (SrcLoc -> TypeArg dim)
-> Maybe SrcLoc -> Maybe (TypeArg dim)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> Maybe SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    unifyTypeArgs TypeArg dim
_ TypeArg dim
_ =
      Maybe (TypeArg dim)
forall a. Maybe a
Nothing
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
uf (Record Map Name (TypeBase dim als)
ts1) (Record Map Name (TypeBase dim als)
ts2)
  | Map Name (TypeBase dim als) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name (TypeBase dim als)
ts1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name (TypeBase dim als) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name (TypeBase dim als)
ts2,
    [Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort (Map Name (TypeBase dim als) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase dim als)
ts1) [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort (Map Name (TypeBase dim als) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase dim als)
ts2) =
    Map Name (TypeBase dim als) -> ScalarTypeBase dim als
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record
      (Map Name (TypeBase dim als) -> ScalarTypeBase dim als)
-> Maybe (Map Name (TypeBase dim als))
-> Maybe (ScalarTypeBase dim als)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((TypeBase dim als, TypeBase dim als) -> Maybe (TypeBase dim als))
-> Map Name (TypeBase dim als, TypeBase dim als)
-> Maybe (Map Name (TypeBase dim als))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
        ((TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als))
-> (TypeBase dim als, TypeBase dim als) -> Maybe (TypeBase dim als)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
uf))
        ((TypeBase dim als
 -> TypeBase dim als -> (TypeBase dim als, TypeBase dim als))
-> Map Name (TypeBase dim als)
-> Map Name (TypeBase dim als)
-> Map Name (TypeBase dim als, TypeBase dim als)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name (TypeBase dim als)
ts1 Map Name (TypeBase dim als)
ts2)
unifyScalarTypes
  Uniqueness -> Uniqueness -> Maybe Uniqueness
uf
  (Arrow als
as1 PName
mn1 TypeBase dim als
t1 (RetType [VName]
dims1 TypeBase dim als
t1'))
  (Arrow als
as2 PName
_ TypeBase dim als
t2 (RetType [VName]
_ TypeBase dim als
t2')) =
    als
-> PName
-> TypeBase dim als
-> RetTypeBase dim als
-> ScalarTypeBase dim als
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow (als
as1 als -> als -> als
forall a. Semigroup a => a -> a -> a
<> als
as2) PName
mn1 (TypeBase dim als -> RetTypeBase dim als -> ScalarTypeBase dim als)
-> Maybe (TypeBase dim als)
-> Maybe (RetTypeBase dim als -> ScalarTypeBase dim als)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU ((Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> Uniqueness -> Uniqueness -> Maybe Uniqueness
forall a b c. (a -> b -> c) -> b -> a -> c
flip Uniqueness -> Uniqueness -> Maybe Uniqueness
uf) TypeBase dim als
t1 TypeBase dim als
t2
      Maybe (RetTypeBase dim als -> ScalarTypeBase dim als)
-> Maybe (RetTypeBase dim als) -> Maybe (ScalarTypeBase dim als)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([VName] -> TypeBase dim als -> RetTypeBase dim als
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims1 (TypeBase dim als -> RetTypeBase dim als)
-> Maybe (TypeBase dim als) -> Maybe (RetTypeBase dim als)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
uf TypeBase dim als
t1' TypeBase dim als
t2')
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
uf (Sum Map Name [TypeBase dim als]
cs1) (Sum Map Name [TypeBase dim als]
cs2)
  | Map Name [TypeBase dim als] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [TypeBase dim als]
cs1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name [TypeBase dim als] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [TypeBase dim als]
cs2,
    [Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort (Map Name [TypeBase dim als] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase dim als]
cs1) [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort (Map Name [TypeBase dim als] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase dim als]
cs2) =
    Map Name [TypeBase dim als] -> ScalarTypeBase dim als
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum
      (Map Name [TypeBase dim als] -> ScalarTypeBase dim als)
-> Maybe (Map Name [TypeBase dim als])
-> Maybe (ScalarTypeBase dim als)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TypeBase dim als], [TypeBase dim als])
 -> Maybe [TypeBase dim als])
-> Map Name ([TypeBase dim als], [TypeBase dim als])
-> Maybe (Map Name [TypeBase dim als])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
        (([TypeBase dim als]
 -> [TypeBase dim als] -> Maybe [TypeBase dim als])
-> ([TypeBase dim als], [TypeBase dim als])
-> Maybe [TypeBase dim als]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als))
-> [TypeBase dim als]
-> [TypeBase dim als]
-> Maybe [TypeBase dim als]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ((Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
uf)))
        (([TypeBase dim als]
 -> [TypeBase dim als] -> ([TypeBase dim als], [TypeBase dim als]))
-> Map Name [TypeBase dim als]
-> Map Name [TypeBase dim als]
-> Map Name ([TypeBase dim als], [TypeBase dim als])
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name [TypeBase dim als]
cs1 Map Name [TypeBase dim als]
cs2)
unifyScalarTypes Uniqueness -> Uniqueness -> Maybe Uniqueness
_ ScalarTypeBase dim als
_ ScalarTypeBase dim als
_ = Maybe (ScalarTypeBase dim als)
forall a. Maybe a
Nothing

-- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal
-- to @y@), meaning @x@ is valid whenever @y@ is.  Ignores sizes.
-- Mostly used for checking uniqueness.
subtypeOf :: TypeBase () () -> TypeBase () () -> Bool
subtypeOf :: TypeBase () () -> TypeBase () () -> Bool
subtypeOf TypeBase () ()
t1 TypeBase () ()
t2 = Maybe (TypeBase () ()) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (TypeBase () ()) -> Bool) -> Maybe (TypeBase () ()) -> Bool
forall a b. (a -> b) -> a -> b
$ (Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase () () -> TypeBase () () -> Maybe (TypeBase () ())
forall als dim.
(Monoid als, ArrayDim dim) =>
(Uniqueness -> Uniqueness -> Maybe Uniqueness)
-> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als)
unifyTypesU Uniqueness -> Uniqueness -> Maybe Uniqueness
unifyUniqueness (TypeBase () () -> TypeBase () ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase () ()
t1) (TypeBase () () -> TypeBase () ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase () ()
t2)
  where
    unifyUniqueness :: Uniqueness -> Uniqueness -> Maybe Uniqueness
unifyUniqueness Uniqueness
u2 Uniqueness
u1 = if Uniqueness
u2 Uniqueness -> Uniqueness -> Bool
`subuniqueOf` Uniqueness
u1 then Uniqueness -> Maybe Uniqueness
forall a. a -> Maybe a
Just Uniqueness
u1 else Maybe Uniqueness
forall a. Maybe a
Nothing

-- | @x `subuniqueOf` y@ is true if @x@ is not less unique than @y@.
subuniqueOf :: Uniqueness -> Uniqueness -> Bool
subuniqueOf :: Uniqueness -> Uniqueness -> Bool
subuniqueOf Uniqueness
Nonunique Uniqueness
Unique = Bool
False
subuniqueOf Uniqueness
_ Uniqueness
_ = Bool
True

-- | Ensure that the dimensions of the RetType are unique by
-- generating new names for them.  This is to avoid name capture.
renameRetType :: MonadTypeChecker m => StructRetType -> m StructRetType
renameRetType :: StructRetType -> m StructRetType
renameRetType (RetType [VName]
dims TypeBase (DimDecl VName) ()
st)
  | [VName]
dims [VName] -> [VName] -> Bool
forall a. Eq a => a -> a -> Bool
/= [VName]
forall a. Monoid a => a
mempty = do
    [VName]
dims' <- (VName -> m VName) -> [VName] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> m VName
forall (m :: * -> *). MonadTypeChecker m => VName -> m VName
newName [VName]
dims
    let m :: Map VName (Subst t)
m = [(VName, Subst t)] -> Map VName (Subst t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Subst t)] -> Map VName (Subst t))
-> [(VName, Subst t)] -> Map VName (Subst t)
forall a b. (a -> b) -> a -> b
$ [VName] -> [Subst t] -> [(VName, Subst t)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
dims ([Subst t] -> [(VName, Subst t)])
-> [Subst t] -> [(VName, Subst t)]
forall a b. (a -> b) -> a -> b
$ (VName -> Subst t) -> [VName] -> [Subst t]
forall a b. (a -> b) -> [a] -> [b]
map (DimDecl VName -> Subst t
forall t. DimDecl VName -> Subst t
SizeSubst (DimDecl VName -> Subst t)
-> (VName -> DimDecl VName) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> (VName -> QualName VName) -> VName -> DimDecl VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName) [VName]
dims'
        st' :: TypeBase (DimDecl VName) ()
st' = TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
forall t. Map VName (Subst t)
m) TypeBase (DimDecl VName) ()
st
    StructRetType -> m StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> m StructRetType)
-> StructRetType -> m StructRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims' TypeBase (DimDecl VName) ()
st'
  | Bool
otherwise =
    StructRetType -> m StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> m StructRetType)
-> StructRetType -> m StructRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase (DimDecl VName) ()
st

evalTypeExp ::
  MonadTypeChecker m =>
  TypeExp Name ->
  m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp :: TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp (TEVar QualName Name
name SrcLoc
loc) = do
  (QualName VName
name', [TypeParam]
ps, StructRetType
t, Liftedness
l) <- SrcLoc
-> QualName Name
-> m (QualName VName, [TypeParam], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc
-> QualName Name
-> m (QualName VName, [TypeParam], StructRetType, Liftedness)
lookupType SrcLoc
loc QualName Name
name
  StructRetType
t' <- StructRetType -> m StructRetType
forall (m :: * -> *).
MonadTypeChecker m =>
StructRetType -> m StructRetType
renameRetType StructRetType
t
  case [TypeParam]
ps of
    [] -> (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> SrcLoc -> TypeExp VName
forall vn. QualName vn -> SrcLoc -> TypeExp vn
TEVar QualName VName
name' SrcLoc
loc, [], StructRetType
t', Liftedness
l)
    [TypeParam]
_ ->
      SrcLoc
-> Notes
-> Doc
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$
        Doc
"Type constructor" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote ([Doc] -> Doc
spread (QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr QualName Name
name Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: (TypeParam -> Doc) -> [TypeParam] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> Doc
forall a. Pretty a => a -> Doc
ppr [TypeParam]
ps))
          Doc -> Doc -> Doc
<+> Doc
"used without any arguments."
--
evalTypeExp (TETuple [TypeExp Name]
ts SrcLoc
loc) = do
  ([TypeExp VName]
ts', [[VName]]
svars, [StructRetType]
ts_s, [Liftedness]
ls) <- [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> ([TypeExp VName], [[VName]], [StructRetType], [Liftedness])
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
unzip4 ([(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> ([TypeExp VName], [[VName]], [StructRetType], [Liftedness]))
-> m [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> m ([TypeExp VName], [[VName]], [StructRetType], [Liftedness])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeExp Name
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> [TypeExp Name]
-> m [(TypeExp VName, [VName], StructRetType, Liftedness)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp [TypeExp Name]
ts
  (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( [TypeExp VName] -> SrcLoc -> TypeExp VName
forall vn. [TypeExp vn] -> SrcLoc -> TypeExp vn
TETuple [TypeExp VName]
ts' SrcLoc
loc,
      [[VName]] -> [VName]
forall a. Monoid a => [a] -> a
mconcat [[VName]]
svars,
      [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ((StructRetType -> [VName]) -> [StructRetType] -> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims [StructRetType]
ts_s) (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ [TypeBase (DimDecl VName) ()] -> ScalarTypeBase (DimDecl VName) ()
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord ([TypeBase (DimDecl VName) ()]
 -> ScalarTypeBase (DimDecl VName) ())
-> [TypeBase (DimDecl VName) ()]
-> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ (StructRetType -> TypeBase (DimDecl VName) ())
-> [StructRetType] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map StructRetType -> TypeBase (DimDecl VName) ()
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType [StructRetType]
ts_s,
      (Liftedness -> Liftedness -> Liftedness)
-> Liftedness -> [Liftedness] -> Liftedness
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Liftedness -> Liftedness -> Liftedness
forall a. Ord a => a -> a -> a
max Liftedness
Unlifted [Liftedness]
ls
    )
--
evalTypeExp t :: TypeExp Name
t@(TERecord [(Name, TypeExp Name)]
fs SrcLoc
loc) = do
  -- Check for duplicate field names.
  let field_names :: [Name]
field_names = ((Name, TypeExp Name) -> Name) -> [(Name, TypeExp Name)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeExp Name) -> Name
forall a b. (a, b) -> a
fst [(Name, TypeExp Name)]
fs
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort [Name]
field_names [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort ([Name] -> [Name]
forall a. Ord a => [a] -> [a]
nubOrd [Name]
field_names)) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    SrcLoc -> Notes -> Doc -> m ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m ()) -> Doc -> m ()
forall a b. (a -> b) -> a -> b
$ Doc
"Duplicate record fields in" Doc -> Doc -> Doc
<+> TypeExp Name -> Doc
forall a. Pretty a => a -> Doc
ppr TypeExp Name
t Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."

  Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
checked <- (TypeExp Name
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> Map Name (TypeExp Name)
-> m (Map Name (TypeExp VName, [VName], StructRetType, Liftedness))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp (Map Name (TypeExp Name)
 -> m (Map
         Name (TypeExp VName, [VName], StructRetType, Liftedness)))
-> Map Name (TypeExp Name)
-> m (Map Name (TypeExp VName, [VName], StructRetType, Liftedness))
forall a b. (a -> b) -> a -> b
$ [(Name, TypeExp Name)] -> Map Name (TypeExp Name)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, TypeExp Name)]
fs
  let fs' :: Map Name (TypeExp VName)
fs' = ((TypeExp VName, [VName], StructRetType, Liftedness)
 -> TypeExp VName)
-> Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
-> Map Name (TypeExp VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(TypeExp VName
x, [VName]
_, StructRetType
_, Liftedness
_) -> TypeExp VName
x) Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
checked
      fs_svars :: [VName]
fs_svars = ((TypeExp VName, [VName], StructRetType, Liftedness) -> [VName])
-> Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
-> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(TypeExp VName
_, [VName]
y, StructRetType
_, Liftedness
_) -> [VName]
y) Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
checked
      ts_s :: Map Name StructRetType
ts_s = ((TypeExp VName, [VName], StructRetType, Liftedness)
 -> StructRetType)
-> Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
-> Map Name StructRetType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(TypeExp VName
_, [VName]
_, StructRetType
z, Liftedness
_) -> StructRetType
z) Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
checked
      ls :: Map Name Liftedness
ls = ((TypeExp VName, [VName], StructRetType, Liftedness) -> Liftedness)
-> Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
-> Map Name Liftedness
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(TypeExp VName
_, [VName]
_, StructRetType
_, Liftedness
v) -> Liftedness
v) Map Name (TypeExp VName, [VName], StructRetType, Liftedness)
checked
  (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( [(Name, TypeExp VName)] -> SrcLoc -> TypeExp VName
forall vn. [(Name, TypeExp vn)] -> SrcLoc -> TypeExp vn
TERecord (Map Name (TypeExp VName) -> [(Name, TypeExp VName)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (TypeExp VName)
fs') SrcLoc
loc,
      [VName]
fs_svars,
      [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ((StructRetType -> [VName]) -> Map Name StructRetType -> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims Map Name StructRetType
ts_s) (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase (DimDecl VName) ())
 -> ScalarTypeBase (DimDecl VName) ())
-> Map Name (TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ (StructRetType -> TypeBase (DimDecl VName) ())
-> Map Name StructRetType -> Map Name (TypeBase (DimDecl VName) ())
forall a b k. (a -> b) -> Map k a -> Map k b
M.map StructRetType -> TypeBase (DimDecl VName) ()
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType Map Name StructRetType
ts_s,
      (Liftedness -> Liftedness -> Liftedness)
-> Liftedness -> Map Name Liftedness -> Liftedness
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Liftedness -> Liftedness -> Liftedness
forall a. Ord a => a -> a -> a
max Liftedness
Unlifted Map Name Liftedness
ls
    )
--
evalTypeExp (TEArray TypeExp Name
t DimExp Name
d SrcLoc
loc) = do
  ([VName]
d_svars, DimExp VName
d', DimDecl VName
d'') <- DimExp Name -> m ([VName], DimExp VName, DimDecl VName)
forall (m :: * -> *).
MonadTypeChecker m =>
DimExp Name -> m ([VName], DimExp VName, DimDecl VName)
checkDimExp DimExp Name
d
  (TypeExp VName
t', [VName]
svars, RetType [VName]
dims TypeBase (DimDecl VName) ()
st, Liftedness
l) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t
  case (Liftedness
l, TypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> Uniqueness
-> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOf TypeBase (DimDecl VName) ()
st ([DimDecl VName] -> ShapeDecl (DimDecl VName)
forall dim. [dim] -> ShapeDecl dim
ShapeDecl [DimDecl VName
d'']) Uniqueness
Nonunique) of
    (Liftedness
Unlifted, TypeBase (DimDecl VName) ()
st') ->
      (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( TypeExp VName -> DimExp VName -> SrcLoc -> TypeExp VName
forall vn. TypeExp vn -> DimExp vn -> SrcLoc -> TypeExp vn
TEArray TypeExp VName
t' DimExp VName
d' SrcLoc
loc,
          [VName]
svars,
          [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
d_svars [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims) TypeBase (DimDecl VName) ()
st',
          Liftedness
Unlifted
        )
    (Liftedness
SizeLifted, TypeBase (DimDecl VName) ()
_) ->
      SrcLoc
-> Notes
-> Doc
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$
        Doc
"Cannot create array with elements of size-lifted type" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (TypeExp Name -> Doc
forall a. Pretty a => a -> Doc
ppr TypeExp Name
t)
          Doc -> Doc -> Doc
<+/> Doc
"(might cause irregular array)."
    (Liftedness
Lifted, TypeBase (DimDecl VName) ()
_) ->
      SrcLoc
-> Notes
-> Doc
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$
        Doc
"Cannot create array with elements of lifted type" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (TypeExp Name -> Doc
forall a. Pretty a => a -> Doc
ppr TypeExp Name
t)
          Doc -> Doc -> Doc
<+/> Doc
"(might contain function)."
  where
    checkDimExp :: DimExp Name -> m ([VName], DimExp VName, DimDecl VName)
checkDimExp DimExp Name
DimExpAny = do
      VName
dv <- Name -> m VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
"d"
      ([VName], DimExp VName, DimDecl VName)
-> m ([VName], DimExp VName, DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName
dv], DimExp VName
forall vn. DimExp vn
DimExpAny, QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
dv)
    checkDimExp (DimExpConst Int
k SrcLoc
dloc) =
      ([VName], DimExp VName, DimDecl VName)
-> m ([VName], DimExp VName, DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Int -> SrcLoc -> DimExp VName
forall vn. Int -> SrcLoc -> DimExp vn
DimExpConst Int
k SrcLoc
dloc, Int -> DimDecl VName
forall vn. Int -> DimDecl vn
ConstDim Int
k)
    checkDimExp (DimExpNamed QualName Name
v SrcLoc
dloc) = do
      QualName VName
v' <- SrcLoc -> QualName Name -> m (QualName VName)
forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m (QualName VName)
checkNamedDim SrcLoc
loc QualName Name
v
      ([VName], DimExp VName, DimDecl VName)
-> m ([VName], DimExp VName, DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], QualName VName -> SrcLoc -> DimExp VName
forall vn. QualName vn -> SrcLoc -> DimExp vn
DimExpNamed QualName VName
v' SrcLoc
dloc, QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v')
--
evalTypeExp (TEUnique TypeExp Name
t SrcLoc
loc) = do
  (TypeExp VName
t', [VName]
svars, RetType [VName]
dims TypeBase (DimDecl VName) ()
st, Liftedness
l) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (TypeBase (DimDecl VName) () -> Bool
forall dim as. TypeBase dim as -> Bool
mayContainArray TypeBase (DimDecl VName) ()
st) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    SrcLoc -> Doc -> m ()
forall (m :: * -> *) loc.
(MonadTypeChecker m, Located loc) =>
loc -> Doc -> m ()
warn SrcLoc
loc (Doc -> m ()) -> Doc -> m ()
forall a b. (a -> b) -> a -> b
$ Doc
"Declaring" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (TypeBase (DimDecl VName) () -> Doc
forall a. Pretty a => a -> Doc
ppr TypeBase (DimDecl VName) ()
st) Doc -> Doc -> Doc
<+> Doc
"as unique has no effect."
  (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeExp VName -> SrcLoc -> TypeExp VName
forall vn. TypeExp vn -> SrcLoc -> TypeExp vn
TEUnique TypeExp VName
t' SrcLoc
loc, [VName]
svars, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) ()
st TypeBase (DimDecl VName) ()
-> Uniqueness -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Unique, Liftedness
l)
  where
    mayContainArray :: TypeBase dim as -> Bool
mayContainArray (Scalar Prim {}) = Bool
False
    mayContainArray Array {} = Bool
True
    mayContainArray (Scalar (Record Map Name (TypeBase dim as)
fs)) = (TypeBase dim as -> Bool) -> Map Name (TypeBase dim as) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeBase dim as -> Bool
mayContainArray Map Name (TypeBase dim as)
fs
    mayContainArray (Scalar TypeVar {}) = Bool
True
    mayContainArray (Scalar Arrow {}) = Bool
False
    mayContainArray (Scalar (Sum Map Name [TypeBase dim as]
cs)) = (([TypeBase dim as] -> Bool) -> Map Name [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (([TypeBase dim as] -> Bool) -> Map Name [TypeBase dim as] -> Bool)
-> ((TypeBase dim as -> Bool) -> [TypeBase dim as] -> Bool)
-> (TypeBase dim as -> Bool)
-> Map Name [TypeBase dim as]
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase dim as -> Bool) -> [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any) TypeBase dim as -> Bool
mayContainArray Map Name [TypeBase dim as]
cs
--
evalTypeExp (TEArrow (Just Name
v) TypeExp Name
t1 TypeExp Name
t2 SrcLoc
loc) = do
  (TypeExp VName
t1', [VName]
svars1, RetType [VName]
dims1 TypeBase (DimDecl VName) ()
st1, Liftedness
_) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t1
  [(Namespace, Name)]
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a.
MonadTypeChecker m =>
[(Namespace, Name)] -> m a -> m a
bindSpaced [(Namespace
Term, Name
v)] (m (TypeExp VName, [VName], StructRetType, Liftedness)
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$ do
    VName
v' <- Namespace -> Name -> SrcLoc -> m VName
forall (m :: * -> *).
MonadTypeChecker m =>
Namespace -> Name -> SrcLoc -> m VName
checkName Namespace
Term Name
v SrcLoc
loc
    VName
-> BoundV
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a.
MonadTypeChecker m =>
VName -> BoundV -> m a -> m a
bindVal VName
v' ([TypeParam] -> TypeBase (DimDecl VName) () -> BoundV
BoundV [] TypeBase (DimDecl VName) ()
st1) (m (TypeExp VName, [VName], StructRetType, Liftedness)
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$ do
      (TypeExp VName
t2', [VName]
svars2, RetType [VName]
dims2 TypeBase (DimDecl VName) ()
st2, Liftedness
_) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t2
      (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( Maybe VName
-> TypeExp VName -> TypeExp VName -> SrcLoc -> TypeExp VName
forall vn.
Maybe vn -> TypeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArrow (VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v') TypeExp VName
t1' TypeExp VName
t2' SrcLoc
loc,
          [VName]
svars1 [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims1 [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
svars2,
          [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> TypeBase (DimDecl VName) ()
-> StructRetType
-> ScalarTypeBase (DimDecl VName) ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow ()
forall a. Monoid a => a
mempty (VName -> PName
Named VName
v') TypeBase (DimDecl VName) ()
st1 ([VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims2 TypeBase (DimDecl VName) ()
st2),
          Liftedness
Lifted
        )
--
evalTypeExp (TEArrow Maybe Name
Nothing TypeExp Name
t1 TypeExp Name
t2 SrcLoc
loc) = do
  (TypeExp VName
t1', [VName]
svars1, RetType [VName]
dims1 TypeBase (DimDecl VName) ()
st1, Liftedness
_) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t1
  (TypeExp VName
t2', [VName]
svars2, RetType [VName]
dims2 TypeBase (DimDecl VName) ()
st2, Liftedness
_) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t2
  (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( Maybe VName
-> TypeExp VName -> TypeExp VName -> SrcLoc -> TypeExp VName
forall vn.
Maybe vn -> TypeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArrow Maybe VName
forall a. Maybe a
Nothing TypeExp VName
t1' TypeExp VName
t2' SrcLoc
loc,
      [VName]
svars1 [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims1 [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
svars2,
      [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> TypeBase (DimDecl VName) ()
-> StructRetType
-> ScalarTypeBase (DimDecl VName) ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow ()
forall a. Monoid a => a
mempty PName
Unnamed TypeBase (DimDecl VName) ()
st1 (StructRetType -> ScalarTypeBase (DimDecl VName) ())
-> StructRetType -> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims2 TypeBase (DimDecl VName) ()
st2,
      Liftedness
Lifted
    )
--
evalTypeExp (TEDim [Name]
dims TypeExp Name
t SrcLoc
loc) = do
  [(Namespace, Name)]
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a.
MonadTypeChecker m =>
[(Namespace, Name)] -> m a -> m a
bindSpaced ((Name -> (Namespace, Name)) -> [Name] -> [(Namespace, Name)]
forall a b. (a -> b) -> [a] -> [b]
map (Namespace
Term,) [Name]
dims) (m (TypeExp VName, [VName], StructRetType, Liftedness)
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$ do
    [VName]
dims' <- (Name -> m VName) -> [Name] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Name -> SrcLoc -> m VName) -> SrcLoc -> Name -> m VName
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Namespace -> Name -> SrcLoc -> m VName
forall (m :: * -> *).
MonadTypeChecker m =>
Namespace -> Name -> SrcLoc -> m VName
checkName Namespace
Term) SrcLoc
loc) [Name]
dims
    [VName]
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a. MonadTypeChecker m => [VName] -> m a -> m a
bindDims [VName]
dims' (m (TypeExp VName, [VName], StructRetType, Liftedness)
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$ do
      (TypeExp VName
t', [VName]
svars, RetType [VName]
t_dims TypeBase (DimDecl VName) ()
st, Liftedness
l) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
t
      (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( [VName] -> TypeExp VName -> SrcLoc -> TypeExp VName
forall vn. [vn] -> TypeExp vn -> SrcLoc -> TypeExp vn
TEDim [VName]
dims' TypeExp VName
t' SrcLoc
loc,
          [VName]
svars,
          [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
dims' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
t_dims) TypeBase (DimDecl VName) ()
st,
          Liftedness -> Liftedness -> Liftedness
forall a. Ord a => a -> a -> a
max Liftedness
l Liftedness
SizeLifted
        )
  where
    bindDims :: [VName] -> m a -> m a
bindDims [] m a
m = m a
m
    bindDims (VName
d : [VName]
ds) m a
m =
      VName -> BoundV -> m a -> m a
forall (m :: * -> *) a.
MonadTypeChecker m =>
VName -> BoundV -> m a -> m a
bindVal VName
d ([TypeParam] -> TypeBase (DimDecl VName) () -> BoundV
BoundV [] (TypeBase (DimDecl VName) () -> BoundV)
-> TypeBase (DimDecl VName) () -> BoundV
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) ())
-> PrimType -> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64) (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ [VName] -> m a -> m a
bindDims [VName]
ds m a
m
--
evalTypeExp t :: TypeExp Name
t@(TESum [(Name, [TypeExp Name])]
cs SrcLoc
loc) = do
  let constructors :: [Name]
constructors = ((Name, [TypeExp Name]) -> Name)
-> [(Name, [TypeExp Name])] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, [TypeExp Name]) -> Name
forall a b. (a, b) -> a
fst [(Name, [TypeExp Name])]
cs
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort [Name]
constructors [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name] -> [Name]
forall a. Ord a => [a] -> [a]
sort ([Name] -> [Name]
forall a. Ord a => [a] -> [a]
nubOrd [Name]
constructors)) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    SrcLoc -> Notes -> Doc -> m ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m ()) -> Doc -> m ()
forall a b. (a -> b) -> a -> b
$ Doc
"Duplicate constructors in" Doc -> Doc -> Doc
<+> TypeExp Name -> Doc
forall a. Pretty a => a -> Doc
ppr TypeExp Name
t

  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Name] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
constructors Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
256) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    SrcLoc -> Notes -> Doc -> m ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty Doc
"Sum types must have less than 256 constructors."

  Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
checked <- (([TypeExp Name]
 -> m [(TypeExp VName, [VName], StructRetType, Liftedness)])
-> Map Name [TypeExp Name]
-> m (Map
        Name [(TypeExp VName, [VName], StructRetType, Liftedness)])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (([TypeExp Name]
  -> m [(TypeExp VName, [VName], StructRetType, Liftedness)])
 -> Map Name [TypeExp Name]
 -> m (Map
         Name [(TypeExp VName, [VName], StructRetType, Liftedness)]))
-> ((TypeExp Name
     -> m (TypeExp VName, [VName], StructRetType, Liftedness))
    -> [TypeExp Name]
    -> m [(TypeExp VName, [VName], StructRetType, Liftedness)])
-> (TypeExp Name
    -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> Map Name [TypeExp Name]
-> m (Map
        Name [(TypeExp VName, [VName], StructRetType, Liftedness)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeExp Name
 -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> [TypeExp Name]
-> m [(TypeExp VName, [VName], StructRetType, Liftedness)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse) TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp (Map Name [TypeExp Name]
 -> m (Map
         Name [(TypeExp VName, [VName], StructRetType, Liftedness)]))
-> Map Name [TypeExp Name]
-> m (Map
        Name [(TypeExp VName, [VName], StructRetType, Liftedness)])
forall a b. (a -> b) -> a -> b
$ [(Name, [TypeExp Name])] -> Map Name [TypeExp Name]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [TypeExp Name])]
cs
  let cs' :: Map Name [TypeExp VName]
cs' = (([(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> [TypeExp VName])
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> Map Name [TypeExp VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([(TypeExp VName, [VName], StructRetType, Liftedness)]
  -> [TypeExp VName])
 -> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> Map Name [TypeExp VName])
-> (((TypeExp VName, [VName], StructRetType, Liftedness)
     -> TypeExp VName)
    -> [(TypeExp VName, [VName], StructRetType, Liftedness)]
    -> [TypeExp VName])
-> ((TypeExp VName, [VName], StructRetType, Liftedness)
    -> TypeExp VName)
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> Map Name [TypeExp VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TypeExp VName, [VName], StructRetType, Liftedness)
 -> TypeExp VName)
-> [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [TypeExp VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (\(TypeExp VName
x, [VName]
_, StructRetType
_, Liftedness
_) -> TypeExp VName
x) Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
checked
      cs_svars :: [VName]
cs_svars = (([(TypeExp VName, [VName], StructRetType, Liftedness)] -> [VName])
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (([(TypeExp VName, [VName], StructRetType, Liftedness)] -> [VName])
 -> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> [VName])
-> (((TypeExp VName, [VName], StructRetType, Liftedness)
     -> [VName])
    -> [(TypeExp VName, [VName], StructRetType, Liftedness)]
    -> [VName])
-> ((TypeExp VName, [VName], StructRetType, Liftedness) -> [VName])
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TypeExp VName, [VName], StructRetType, Liftedness) -> [VName])
-> [(TypeExp VName, [VName], StructRetType, Liftedness)] -> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap) (\(TypeExp VName
_, [VName]
y, StructRetType
_, Liftedness
_) -> [VName]
y) Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
checked
      ts_s :: Map Name [StructRetType]
ts_s = (([(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> [StructRetType])
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> Map Name [StructRetType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([(TypeExp VName, [VName], StructRetType, Liftedness)]
  -> [StructRetType])
 -> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> Map Name [StructRetType])
-> (((TypeExp VName, [VName], StructRetType, Liftedness)
     -> StructRetType)
    -> [(TypeExp VName, [VName], StructRetType, Liftedness)]
    -> [StructRetType])
-> ((TypeExp VName, [VName], StructRetType, Liftedness)
    -> StructRetType)
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> Map Name [StructRetType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TypeExp VName, [VName], StructRetType, Liftedness)
 -> StructRetType)
-> [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [StructRetType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (\(TypeExp VName
_, [VName]
_, StructRetType
z, Liftedness
_) -> StructRetType
z) Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
checked
      ls :: [Liftedness]
ls = (([(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> [Liftedness])
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [Liftedness]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (([(TypeExp VName, [VName], StructRetType, Liftedness)]
  -> [Liftedness])
 -> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
 -> [Liftedness])
-> (((TypeExp VName, [VName], StructRetType, Liftedness)
     -> Liftedness)
    -> [(TypeExp VName, [VName], StructRetType, Liftedness)]
    -> [Liftedness])
-> ((TypeExp VName, [VName], StructRetType, Liftedness)
    -> Liftedness)
-> Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [Liftedness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TypeExp VName, [VName], StructRetType, Liftedness) -> Liftedness)
-> [(TypeExp VName, [VName], StructRetType, Liftedness)]
-> [Liftedness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (\(TypeExp VName
_, [VName]
_, StructRetType
_, Liftedness
v) -> Liftedness
v) Map Name [(TypeExp VName, [VName], StructRetType, Liftedness)]
checked
  (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( [(Name, [TypeExp VName])] -> SrcLoc -> TypeExp VName
forall vn. [(Name, [TypeExp vn])] -> SrcLoc -> TypeExp vn
TESum (Map Name [TypeExp VName] -> [(Name, [TypeExp VName])]
forall k a. Map k a -> [(k, a)]
M.toList Map Name [TypeExp VName]
cs') SrcLoc
loc,
      [VName]
cs_svars,
      [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (([StructRetType] -> [VName]) -> Map Name [StructRetType] -> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((StructRetType -> [VName]) -> [StructRetType] -> [VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims) Map Name [StructRetType]
ts_s) (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$
        ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase (DimDecl VName) ()]
-> ScalarTypeBase (DimDecl VName) ()
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase (DimDecl VName) ()]
 -> ScalarTypeBase (DimDecl VName) ())
-> Map Name [TypeBase (DimDecl VName) ()]
-> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ ([StructRetType] -> [TypeBase (DimDecl VName) ()])
-> Map Name [StructRetType]
-> Map Name [TypeBase (DimDecl VName) ()]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructRetType -> TypeBase (DimDecl VName) ())
-> [StructRetType] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map StructRetType -> TypeBase (DimDecl VName) ()
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType) Map Name [StructRetType]
ts_s,
      (Liftedness -> Liftedness -> Liftedness)
-> Liftedness -> [Liftedness] -> Liftedness
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Liftedness -> Liftedness -> Liftedness
forall a. Ord a => a -> a -> a
max Liftedness
Unlifted [Liftedness]
ls
    )
evalTypeExp ote :: TypeExp Name
ote@TEApply {} = do
  (QualName Name
tname, SrcLoc
tname_loc, [TypeArgExp Name]
targs) <- TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
rootAndArgs TypeExp Name
ote
  (QualName VName
tname', [TypeParam]
ps, StructRetType
tname_t, Liftedness
l) <- SrcLoc
-> QualName Name
-> m (QualName VName, [TypeParam], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc
-> QualName Name
-> m (QualName VName, [TypeParam], StructRetType, Liftedness)
lookupType SrcLoc
tloc QualName Name
tname
  RetType [VName]
t_dims TypeBase (DimDecl VName) ()
t <- StructRetType -> m StructRetType
forall (m :: * -> *).
MonadTypeChecker m =>
StructRetType -> m StructRetType
renameRetType StructRetType
tname_t
  if [TypeParam] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeParam]
ps Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [TypeArgExp Name] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeArgExp Name]
targs
    then
      SrcLoc
-> Notes
-> Doc
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
tloc Notes
forall a. Monoid a => a
mempty (Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness))
-> Doc -> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall a b. (a -> b) -> a -> b
$
        Doc
"Type constructor" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (QualName Name -> Doc
forall a. Pretty a => a -> Doc
ppr QualName Name
tname) Doc -> Doc -> Doc
<+> Doc
"requires" Doc -> Doc -> Doc
<+> Int -> Doc
forall a. Pretty a => a -> Doc
ppr ([TypeParam] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeParam]
ps)
          Doc -> Doc -> Doc
<+> Doc
"arguments, but provided"
          Doc -> Doc -> Doc
<+> Int -> Doc
forall a. Pretty a => a -> Doc
ppr ([TypeArgExp Name] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeArgExp Name]
targs) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."
    else do
      ([TypeArgExp VName]
targs', [[VName]]
dims, [Map VName (Subst StructRetType)]
substs) <- [(TypeArgExp VName, [VName], Map VName (Subst StructRetType))]
-> ([TypeArgExp VName], [[VName]],
    [Map VName (Subst StructRetType)])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(TypeArgExp VName, [VName], Map VName (Subst StructRetType))]
 -> ([TypeArgExp VName], [[VName]],
     [Map VName (Subst StructRetType)]))
-> m [(TypeArgExp VName, [VName], Map VName (Subst StructRetType))]
-> m ([TypeArgExp VName], [[VName]],
      [Map VName (Subst StructRetType)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeParam
 -> TypeArgExp Name
 -> m (TypeArgExp VName, [VName], Map VName (Subst StructRetType)))
-> [TypeParam]
-> [TypeArgExp Name]
-> m [(TypeArgExp VName, [VName], Map VName (Subst StructRetType))]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeParam
-> TypeArgExp Name
-> m (TypeArgExp VName, [VName], Map VName (Subst StructRetType))
forall (m :: * -> *) k.
(MonadTypeChecker m, Eq k, IsName k) =>
TypeParamBase k
-> TypeArgExp Name
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
checkArgApply [TypeParam]
ps [TypeArgExp Name]
targs
      (TypeExp VName, [VName], StructRetType, Liftedness)
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( (TypeExp VName -> TypeArgExp VName -> TypeExp VName)
-> TypeExp VName -> [TypeArgExp VName] -> TypeExp VName
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\TypeExp VName
x TypeArgExp VName
y -> TypeExp VName -> TypeArgExp VName -> SrcLoc -> TypeExp VName
forall vn. TypeExp vn -> TypeArgExp vn -> SrcLoc -> TypeExp vn
TEApply TypeExp VName
x TypeArgExp VName
y SrcLoc
tloc) (QualName VName -> SrcLoc -> TypeExp VName
forall vn. QualName vn -> SrcLoc -> TypeExp vn
TEVar QualName VName
tname' SrcLoc
tname_loc) [TypeArgExp VName]
targs',
          [],
          [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
t_dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [[VName]] -> [VName]
forall a. Monoid a => [a] -> a
mconcat [[VName]]
dims) (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` [Map VName (Subst StructRetType)]
-> Map VName (Subst StructRetType)
forall a. Monoid a => [a] -> a
mconcat [Map VName (Subst StructRetType)]
substs) TypeBase (DimDecl VName) ()
t,
          Liftedness
l
        )
  where
    tloc :: SrcLoc
tloc = TypeExp Name -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeExp Name
ote

    rootAndArgs :: MonadTypeChecker m => TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
    rootAndArgs :: TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
rootAndArgs (TEVar QualName Name
qn SrcLoc
loc) = (QualName Name, SrcLoc, [TypeArgExp Name])
-> m (QualName Name, SrcLoc, [TypeArgExp Name])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName Name
qn, SrcLoc
loc, [])
    rootAndArgs (TEApply TypeExp Name
op TypeArgExp Name
arg SrcLoc
_) = do
      (QualName Name
op', SrcLoc
loc, [TypeArgExp Name]
args) <- TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])
rootAndArgs TypeExp Name
op
      (QualName Name, SrcLoc, [TypeArgExp Name])
-> m (QualName Name, SrcLoc, [TypeArgExp Name])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName Name
op', SrcLoc
loc, [TypeArgExp Name]
args [TypeArgExp Name] -> [TypeArgExp Name] -> [TypeArgExp Name]
forall a. [a] -> [a] -> [a]
++ [TypeArgExp Name
arg])
    rootAndArgs TypeExp Name
te' =
      SrcLoc
-> Notes -> Doc -> m (QualName Name, SrcLoc, [TypeArgExp Name])
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError (TypeExp Name -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeExp Name
te') Notes
forall a. Monoid a => a
mempty (Doc -> m (QualName Name, SrcLoc, [TypeArgExp Name]))
-> Doc -> m (QualName Name, SrcLoc, [TypeArgExp Name])
forall a b. (a -> b) -> a -> b
$
        Doc
"Type" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (TypeExp Name -> Doc
forall a. Pretty a => a -> Doc
ppr TypeExp Name
te') Doc -> Doc -> Doc
<+> Doc
"is not a type constructor."

    checkArgApply :: TypeParamBase k
-> TypeArgExp Name
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
checkArgApply (TypeParamDim k
pv SrcLoc
_) (TypeArgExpDim (DimExpNamed QualName Name
v SrcLoc
dloc) SrcLoc
loc) = do
      QualName VName
v' <- SrcLoc -> QualName Name -> m (QualName VName)
forall (m :: * -> *).
MonadTypeChecker m =>
SrcLoc -> QualName Name -> m (QualName VName)
checkNamedDim SrcLoc
loc QualName Name
v
      (TypeArgExp VName, [VName], Map k (Subst StructRetType))
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( DimExp VName -> SrcLoc -> TypeArgExp VName
forall vn. DimExp vn -> SrcLoc -> TypeArgExp vn
TypeArgExpDim (QualName VName -> SrcLoc -> DimExp VName
forall vn. QualName vn -> SrcLoc -> DimExp vn
DimExpNamed QualName VName
v' SrcLoc
dloc) SrcLoc
loc,
          [],
          k -> Subst StructRetType -> Map k (Subst StructRetType)
forall k a. k -> a -> Map k a
M.singleton k
pv (Subst StructRetType -> Map k (Subst StructRetType))
-> Subst StructRetType -> Map k (Subst StructRetType)
forall a b. (a -> b) -> a -> b
$ DimDecl VName -> Subst StructRetType
forall t. DimDecl VName -> Subst t
SizeSubst (DimDecl VName -> Subst StructRetType)
-> DimDecl VName -> Subst StructRetType
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v'
        )
    checkArgApply (TypeParamDim k
pv SrcLoc
_) (TypeArgExpDim (DimExpConst Int
x SrcLoc
dloc) SrcLoc
loc) =
      (TypeArgExp VName, [VName], Map k (Subst StructRetType))
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( DimExp VName -> SrcLoc -> TypeArgExp VName
forall vn. DimExp vn -> SrcLoc -> TypeArgExp vn
TypeArgExpDim (Int -> SrcLoc -> DimExp VName
forall vn. Int -> SrcLoc -> DimExp vn
DimExpConst Int
x SrcLoc
dloc) SrcLoc
loc,
          [],
          k -> Subst StructRetType -> Map k (Subst StructRetType)
forall k a. k -> a -> Map k a
M.singleton k
pv (Subst StructRetType -> Map k (Subst StructRetType))
-> Subst StructRetType -> Map k (Subst StructRetType)
forall a b. (a -> b) -> a -> b
$ DimDecl VName -> Subst StructRetType
forall t. DimDecl VName -> Subst t
SizeSubst (DimDecl VName -> Subst StructRetType)
-> DimDecl VName -> Subst StructRetType
forall a b. (a -> b) -> a -> b
$ Int -> DimDecl VName
forall vn. Int -> DimDecl vn
ConstDim Int
x
        )
    checkArgApply (TypeParamDim k
pv SrcLoc
_) (TypeArgExpDim DimExp Name
DimExpAny SrcLoc
loc) = do
      VName
d <- Name -> m VName
forall (m :: * -> *). MonadTypeChecker m => Name -> m VName
newTypeName Name
"d"
      (TypeArgExp VName, [VName], Map k (Subst StructRetType))
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( DimExp VName -> SrcLoc -> TypeArgExp VName
forall vn. DimExp vn -> SrcLoc -> TypeArgExp vn
TypeArgExpDim DimExp VName
forall vn. DimExp vn
DimExpAny SrcLoc
loc,
          [VName
d],
          k -> Subst StructRetType -> Map k (Subst StructRetType)
forall k a. k -> a -> Map k a
M.singleton k
pv (Subst StructRetType -> Map k (Subst StructRetType))
-> Subst StructRetType -> Map k (Subst StructRetType)
forall a b. (a -> b) -> a -> b
$ DimDecl VName -> Subst StructRetType
forall t. DimDecl VName -> Subst t
SizeSubst (DimDecl VName -> Subst StructRetType)
-> DimDecl VName -> Subst StructRetType
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d
        )
    checkArgApply (TypeParamType Liftedness
_ k
pv SrcLoc
_) (TypeArgExpType TypeExp Name
te) = do
      (TypeExp VName
te', [VName]
svars, RetType [VName]
dims TypeBase (DimDecl VName) ()
st, Liftedness
_) <- TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
te
      (TypeArgExp VName, [VName], Map k (Subst StructRetType))
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( TypeExp VName -> TypeArgExp VName
forall vn. TypeExp vn -> TypeArgExp vn
TypeArgExpType TypeExp VName
te',
          [VName]
svars [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
dims,
          k -> Subst StructRetType -> Map k (Subst StructRetType)
forall k a. k -> a -> Map k a
M.singleton k
pv (Subst StructRetType -> Map k (Subst StructRetType))
-> Subst StructRetType -> Map k (Subst StructRetType)
forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructRetType -> Subst StructRetType
forall t. [TypeParam] -> t -> Subst t
Subst [] (StructRetType -> Subst StructRetType)
-> StructRetType -> Subst StructRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase (DimDecl VName) ()
st
        )
    checkArgApply TypeParamBase k
p TypeArgExp Name
a =
      SrcLoc
-> Notes
-> Doc
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
tloc Notes
forall a. Monoid a => a
mempty (Doc -> m (TypeArgExp VName, [VName], Map k (Subst StructRetType)))
-> Doc
-> m (TypeArgExp VName, [VName], Map k (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$
        Doc
"Type argument" Doc -> Doc -> Doc
<+> TypeArgExp Name -> Doc
forall a. Pretty a => a -> Doc
ppr TypeArgExp Name
a
          Doc -> Doc -> Doc
<+> Doc
"not valid for a type parameter"
          Doc -> Doc -> Doc
<+> TypeParamBase k -> Doc
forall a. Pretty a => a -> Doc
ppr TypeParamBase k
p Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."

checkTypeExp ::
  MonadTypeChecker m =>
  TypeExp Name ->
  m (TypeExp VName, [VName], StructRetType, Liftedness)
checkTypeExp :: TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
checkTypeExp TypeExp Name
te = do
  TypeExp Name -> m ()
forall (m :: * -> *). MonadTypeChecker m => TypeExp Name -> m ()
checkForDuplicateNamesInType TypeExp Name
te
  TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
forall (m :: * -> *).
MonadTypeChecker m =>
TypeExp Name
-> m (TypeExp VName, [VName], StructRetType, Liftedness)
evalTypeExp TypeExp Name
te

-- | Check for duplication of names inside a pattern group.  Produces
-- a description of all names used in the pattern group.
checkForDuplicateNames ::
  MonadTypeChecker m =>
  [UncheckedPat] ->
  m ()
checkForDuplicateNames :: [UncheckedPat] -> m ()
checkForDuplicateNames = (StateT (Map Name SrcLoc) m () -> Map Name SrcLoc -> m ()
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
`evalStateT` Map Name SrcLoc
forall a. Monoid a => a
mempty) (StateT (Map Name SrcLoc) m () -> m ())
-> ([UncheckedPat] -> StateT (Map Name SrcLoc) m ())
-> [UncheckedPat]
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UncheckedPat -> StateT (Map Name SrcLoc) m ())
-> [UncheckedPat] -> StateT (Map Name SrcLoc) m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ UncheckedPat -> StateT (Map Name SrcLoc) m ()
forall k (t :: (* -> *) -> * -> *) (m :: * -> *) (f :: * -> *).
(MonadState (Map k SrcLoc) (t m), Pretty k, MonadTypeChecker m,
 MonadTrans t, Ord k) =>
PatBase f k -> t m ()
check
  where
    check :: PatBase f k -> t m ()
check (Id k
v f PatType
_ SrcLoc
loc) = k -> SrcLoc -> t m ()
forall k a (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadState (Map k a) (t m), Located a, Pretty k,
 MonadTypeChecker m, MonadTrans t, Ord k) =>
k -> a -> t m ()
seen k
v SrcLoc
loc
    check (PatParens PatBase f k
p SrcLoc
_) = PatBase f k -> t m ()
check PatBase f k
p
    check (PatAttr AttrInfo k
_ PatBase f k
p SrcLoc
_) = PatBase f k -> t m ()
check PatBase f k
p
    check Wildcard {} = () -> t m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    check (TuplePat [PatBase f k]
ps SrcLoc
_) = (PatBase f k -> t m ()) -> [PatBase f k] -> t m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ PatBase f k -> t m ()
check [PatBase f k]
ps
    check (RecordPat [(Name, PatBase f k)]
fs SrcLoc
_) = ((Name, PatBase f k) -> t m ()) -> [(Name, PatBase f k)] -> t m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (PatBase f k -> t m ()
check (PatBase f k -> t m ())
-> ((Name, PatBase f k) -> PatBase f k)
-> (Name, PatBase f k)
-> t m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatBase f k) -> PatBase f k
forall a b. (a, b) -> b
snd) [(Name, PatBase f k)]
fs
    check (PatAscription PatBase f k
p TypeDeclBase f k
_ SrcLoc
_) = PatBase f k -> t m ()
check PatBase f k
p
    check PatLit {} = () -> t m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    check (PatConstr Name
_ f PatType
_ [PatBase f k]
ps SrcLoc
_) = (PatBase f k -> t m ()) -> [PatBase f k] -> t m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ PatBase f k -> t m ()
check [PatBase f k]
ps

    seen :: k -> a -> t m ()
seen k
v a
loc = do
      Maybe a
already <- (Map k a -> Maybe a) -> t m (Maybe a)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Map k a -> Maybe a) -> t m (Maybe a))
-> (Map k a -> Maybe a) -> t m (Maybe a)
forall a b. (a -> b) -> a -> b
$ k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
v
      case Maybe a
already of
        Just a
prev_loc ->
          m () -> t m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> t m ()) -> m () -> t m ()
forall a b. (a -> b) -> a -> b
$
            a -> Notes -> Doc -> m ()
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError a
loc Notes
forall a. Monoid a => a
mempty (Doc -> m ()) -> Doc -> m ()
forall a b. (a -> b) -> a -> b
$
              Doc
"Name" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (k -> Doc
forall a. Pretty a => a -> Doc
ppr k
v) Doc -> Doc -> Doc
<+> Doc
"also bound at"
                Doc -> Doc -> Doc
<+> String -> Doc
text (a -> String
forall a. Located a => a -> String
locStr a
prev_loc) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."
        Maybe a
Nothing ->
          (Map k a -> Map k a) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map k a -> Map k a) -> t m ()) -> (Map k a -> Map k a) -> t m ()
forall a b. (a -> b) -> a -> b
$ k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
v a
loc

-- | Check whether the type contains arrow types that define the same
-- parameter.  These might also exist further down, but that's not
-- really a problem - we mostly do this checking to help the user,
-- since it is likely an error, but it's easy to assign a semantics to
-- it (normal name shadowing).
checkForDuplicateNamesInType ::
  MonadTypeChecker m =>
  TypeExp Name ->
  m ()
checkForDuplicateNamesInType :: TypeExp Name -> m ()
checkForDuplicateNamesInType = Map Name SrcLoc -> TypeExp Name -> m ()
forall vn (m :: * -> *).
(MonadTypeChecker m, Pretty vn, Ord vn) =>
Map vn SrcLoc -> TypeExp vn -> m ()
check Map Name SrcLoc
forall a. Monoid a => a
mempty
  where
    bad :: a -> loc -> a -> m a
bad a
v loc
loc a
prev_loc =
      loc -> Notes -> Doc -> m a
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError loc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m a) -> Doc -> m a
forall a b. (a -> b) -> a -> b
$
        String -> Doc
text String
"Name" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (a -> Doc
forall a. Pretty a => a -> Doc
ppr a
v)
          Doc -> Doc -> Doc
<+> Doc
"also bound at"
          Doc -> Doc -> Doc
<+> String -> Doc
text (a -> String
forall a. Located a => a -> String
locStr a
prev_loc) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."

    check :: Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen (TEArrow (Just vn
v) TypeExp vn
t1 TypeExp vn
t2 SrcLoc
loc)
      | Just SrcLoc
prev_loc <- vn -> Map vn SrcLoc -> Maybe SrcLoc
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup vn
v Map vn SrcLoc
seen =
        vn -> SrcLoc -> SrcLoc -> m ()
forall (m :: * -> *) a loc a a.
(MonadTypeChecker m, Pretty a, Located loc, Located a) =>
a -> loc -> a -> m a
bad vn
v SrcLoc
loc SrcLoc
prev_loc
      | Bool
otherwise =
        Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen' TypeExp vn
t1 m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen' TypeExp vn
t2
      where
        seen' :: Map vn SrcLoc
seen' = vn -> SrcLoc -> Map vn SrcLoc -> Map vn SrcLoc
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert vn
v SrcLoc
loc Map vn SrcLoc
seen
    check Map vn SrcLoc
seen (TEArrow Maybe vn
Nothing TypeExp vn
t1 TypeExp vn
t2 SrcLoc
_) =
      Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t1 m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t2
    check Map vn SrcLoc
seen (TETuple [TypeExp vn]
ts SrcLoc
_) = (TypeExp vn -> m ()) -> [TypeExp vn] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen) [TypeExp vn]
ts
    check Map vn SrcLoc
seen (TERecord [(Name, TypeExp vn)]
fs SrcLoc
_) = ((Name, TypeExp vn) -> m ()) -> [(Name, TypeExp vn)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen (TypeExp vn -> m ())
-> ((Name, TypeExp vn) -> TypeExp vn) -> (Name, TypeExp vn) -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeExp vn) -> TypeExp vn
forall a b. (a, b) -> b
snd) [(Name, TypeExp vn)]
fs
    check Map vn SrcLoc
seen (TEUnique TypeExp vn
t SrcLoc
_) = Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t
    check Map vn SrcLoc
seen (TESum [(Name, [TypeExp vn])]
cs SrcLoc
_) = ((Name, [TypeExp vn]) -> m [()]) -> [(Name, [TypeExp vn])] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((TypeExp vn -> m ()) -> [TypeExp vn] -> m [()]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen) ([TypeExp vn] -> m [()])
-> ((Name, [TypeExp vn]) -> [TypeExp vn])
-> (Name, [TypeExp vn])
-> m [()]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [TypeExp vn]) -> [TypeExp vn]
forall a b. (a, b) -> b
snd) [(Name, [TypeExp vn])]
cs
    check Map vn SrcLoc
seen (TEApply TypeExp vn
t1 (TypeArgExpType TypeExp vn
t2) SrcLoc
_) =
      Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t1 m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t2
    check Map vn SrcLoc
seen (TEApply TypeExp vn
t1 TypeArgExpDim {} SrcLoc
_) =
      Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t1
    check Map vn SrcLoc
seen (TEDim (vn
v : [vn]
vs) TypeExp vn
t SrcLoc
loc)
      | Just SrcLoc
prev_loc <- vn -> Map vn SrcLoc -> Maybe SrcLoc
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup vn
v Map vn SrcLoc
seen =
        vn -> SrcLoc -> SrcLoc -> m ()
forall (m :: * -> *) a loc a a.
(MonadTypeChecker m, Pretty a, Located loc, Located a) =>
a -> loc -> a -> m a
bad vn
v SrcLoc
loc SrcLoc
prev_loc
      | Bool
otherwise =
        Map vn SrcLoc -> TypeExp vn -> m ()
check (vn -> SrcLoc -> Map vn SrcLoc -> Map vn SrcLoc
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert vn
v SrcLoc
loc Map vn SrcLoc
seen) ([vn] -> TypeExp vn -> SrcLoc -> TypeExp vn
forall vn. [vn] -> TypeExp vn -> SrcLoc -> TypeExp vn
TEDim [vn]
vs TypeExp vn
t SrcLoc
loc)
    check Map vn SrcLoc
seen (TEDim [] TypeExp vn
t SrcLoc
_) =
      Map vn SrcLoc -> TypeExp vn -> m ()
check Map vn SrcLoc
seen TypeExp vn
t
    check Map vn SrcLoc
_ TEArray {} = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    check Map vn SrcLoc
_ TEVar {} = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | @checkTypeParams ps m@ checks the type parameters @ps@, then
-- invokes the continuation @m@ with the checked parameters, while
-- extending the monadic name map with @ps@.
checkTypeParams ::
  MonadTypeChecker m =>
  [TypeParamBase Name] ->
  ([TypeParamBase VName] -> m a) ->
  m a
checkTypeParams :: [TypeParamBase Name] -> ([TypeParam] -> m a) -> m a
checkTypeParams [TypeParamBase Name]
ps [TypeParam] -> m a
m =
  [(Namespace, Name)] -> m a -> m a
forall (m :: * -> *) a.
MonadTypeChecker m =>
[(Namespace, Name)] -> m a -> m a
bindSpaced ((TypeParamBase Name -> (Namespace, Name))
-> [TypeParamBase Name] -> [(Namespace, Name)]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase Name -> (Namespace, Name)
forall b. TypeParamBase b -> (Namespace, b)
typeParamSpace [TypeParamBase Name]
ps) (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$
    [TypeParam] -> m a
m ([TypeParam] -> m a) -> m [TypeParam] -> m a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< StateT (Map (Namespace, Name) SrcLoc) m [TypeParam]
-> Map (Namespace, Name) SrcLoc -> m [TypeParam]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT ((TypeParamBase Name
 -> StateT (Map (Namespace, Name) SrcLoc) m TypeParam)
-> [TypeParamBase Name]
-> StateT (Map (Namespace, Name) SrcLoc) m [TypeParam]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeParamBase Name
-> StateT (Map (Namespace, Name) SrcLoc) m TypeParam
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadState (Map (Namespace, Name) SrcLoc) (t m),
 MonadTypeChecker m, MonadTrans t) =>
TypeParamBase Name -> t m TypeParam
checkTypeParam [TypeParamBase Name]
ps) Map (Namespace, Name) SrcLoc
forall a. Monoid a => a
mempty
  where
    typeParamSpace :: TypeParamBase b -> (Namespace, b)
typeParamSpace (TypeParamDim b
pv SrcLoc
_) = (Namespace
Term, b
pv)
    typeParamSpace (TypeParamType Liftedness
_ b
pv SrcLoc
_) = (Namespace
Type, b
pv)

    checkParamName :: Namespace -> Name -> SrcLoc -> t m VName
checkParamName Namespace
ns Name
v SrcLoc
loc = do
      Maybe SrcLoc
seen <- (Map (Namespace, Name) SrcLoc -> Maybe SrcLoc)
-> t m (Maybe SrcLoc)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Map (Namespace, Name) SrcLoc -> Maybe SrcLoc)
 -> t m (Maybe SrcLoc))
-> (Map (Namespace, Name) SrcLoc -> Maybe SrcLoc)
-> t m (Maybe SrcLoc)
forall a b. (a -> b) -> a -> b
$ (Namespace, Name) -> Map (Namespace, Name) SrcLoc -> Maybe SrcLoc
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Namespace
ns, Name
v)
      case Maybe SrcLoc
seen of
        Just SrcLoc
prev ->
          m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$
            SrcLoc -> Notes -> Doc -> m VName
forall (m :: * -> *) loc a.
(MonadTypeChecker m, Located loc) =>
loc -> Notes -> Doc -> m a
typeError SrcLoc
loc Notes
forall a. Monoid a => a
mempty (Doc -> m VName) -> Doc -> m VName
forall a b. (a -> b) -> a -> b
$
              String -> Doc
text String
"Type parameter" Doc -> Doc -> Doc
<+> Doc -> Doc
pquote (Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
v)
                Doc -> Doc -> Doc
<+> Doc
"previously defined at"
                Doc -> Doc -> Doc
<+> String -> Doc
text (SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
prev) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"."
        Maybe SrcLoc
Nothing -> do
          (Map (Namespace, Name) SrcLoc -> Map (Namespace, Name) SrcLoc)
-> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map (Namespace, Name) SrcLoc -> Map (Namespace, Name) SrcLoc)
 -> t m ())
-> (Map (Namespace, Name) SrcLoc -> Map (Namespace, Name) SrcLoc)
-> t m ()
forall a b. (a -> b) -> a -> b
$ (Namespace, Name)
-> SrcLoc
-> Map (Namespace, Name) SrcLoc
-> Map (Namespace, Name) SrcLoc
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (Namespace
ns, Name
v) SrcLoc
loc
          m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ Namespace -> Name -> SrcLoc -> m VName
forall (m :: * -> *).
MonadTypeChecker m =>
Namespace -> Name -> SrcLoc -> m VName
checkName Namespace
ns Name
v SrcLoc
loc

    checkTypeParam :: TypeParamBase Name -> t m TypeParam
checkTypeParam (TypeParamDim Name
pv SrcLoc
loc) =
      VName -> SrcLoc -> TypeParam
forall vn. vn -> SrcLoc -> TypeParamBase vn
TypeParamDim (VName -> SrcLoc -> TypeParam)
-> t m VName -> t m (SrcLoc -> TypeParam)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Namespace -> Name -> SrcLoc -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadState (Map (Namespace, Name) SrcLoc) (t m),
 MonadTypeChecker m, MonadTrans t) =>
Namespace -> Name -> SrcLoc -> t m VName
checkParamName Namespace
Term Name
pv SrcLoc
loc t m (SrcLoc -> TypeParam) -> t m SrcLoc -> t m TypeParam
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    checkTypeParam (TypeParamType Liftedness
l Name
pv SrcLoc
loc) =
      Liftedness -> VName -> SrcLoc -> TypeParam
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
l (VName -> SrcLoc -> TypeParam)
-> t m VName -> t m (SrcLoc -> TypeParam)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Namespace -> Name -> SrcLoc -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadState (Map (Namespace, Name) SrcLoc) (t m),
 MonadTypeChecker m, MonadTrans t) =>
Namespace -> Name -> SrcLoc -> t m VName
checkParamName Namespace
Type Name
pv SrcLoc
loc t m (SrcLoc -> TypeParam) -> t m SrcLoc -> t m TypeParam
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> t m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

-- | Construct a type argument corresponding to a type parameter.
typeParamToArg :: TypeParam -> StructTypeArg
typeParamToArg :: TypeParam -> StructTypeArg
typeParamToArg (TypeParamDim VName
v SrcLoc
ploc) =
  DimDecl VName -> SrcLoc -> StructTypeArg
forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim (QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) SrcLoc
ploc
typeParamToArg (TypeParamType Liftedness
_ VName
v SrcLoc
ploc) =
  TypeBase (DimDecl VName) () -> SrcLoc -> StructTypeArg
forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ ()
-> Uniqueness
-> TypeName
-> [StructTypeArg]
-> ScalarTypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> TypeName
typeName VName
v) []) SrcLoc
ploc

-- | A type substitution may be a substitution or a yet-unknown
-- substitution (but which is certainly an overloaded primitive
-- type!).  The latter is used to remove aliases from types that are
-- yet-unknown but that we know cannot carry aliases (see issue #682).
data Subst t = Subst [TypeParam] t | PrimSubst | SizeSubst (DimDecl VName)
  deriving (Int -> Subst t -> ShowS
[Subst t] -> ShowS
Subst t -> String
(Int -> Subst t -> ShowS)
-> (Subst t -> String) -> ([Subst t] -> ShowS) -> Show (Subst t)
forall t. Show t => Int -> Subst t -> ShowS
forall t. Show t => [Subst t] -> ShowS
forall t. Show t => Subst t -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Subst t] -> ShowS
$cshowList :: forall t. Show t => [Subst t] -> ShowS
show :: Subst t -> String
$cshow :: forall t. Show t => Subst t -> String
showsPrec :: Int -> Subst t -> ShowS
$cshowsPrec :: forall t. Show t => Int -> Subst t -> ShowS
Show)

instance Pretty t => Pretty (Subst t) where
  ppr :: Subst t -> Doc
ppr (Subst [] t
t) = t -> Doc
forall a. Pretty a => a -> Doc
ppr t
t
  ppr (Subst [TypeParam]
tps t
t) = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((TypeParam -> Doc) -> [TypeParam] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> Doc
forall a. Pretty a => a -> Doc
ppr [TypeParam]
tps) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
colon Doc -> Doc -> Doc
<+> t -> Doc
forall a. Pretty a => a -> Doc
ppr t
t
  ppr Subst t
PrimSubst = Doc
"#<primsubst>"
  ppr (SizeSubst DimDecl VName
d) = DimDecl VName -> Doc
forall a. Pretty a => a -> Doc
ppr DimDecl VName
d

-- | Create a type substitution corresponding to a type binding.
substFromAbbr :: TypeBinding -> Subst StructRetType
substFromAbbr :: TypeBinding -> Subst StructRetType
substFromAbbr (TypeAbbr Liftedness
_ [TypeParam]
ps StructRetType
rt) = [TypeParam] -> StructRetType -> Subst StructRetType
forall t. [TypeParam] -> t -> Subst t
Subst [TypeParam]
ps StructRetType
rt

-- | Substitutions to apply in a type.
type TypeSubs = VName -> Maybe (Subst StructRetType)

instance Functor Subst where
  fmap :: (a -> b) -> Subst a -> Subst b
fmap a -> b
f (Subst [TypeParam]
ps a
t) = [TypeParam] -> b -> Subst b
forall t. [TypeParam] -> t -> Subst t
Subst [TypeParam]
ps (b -> Subst b) -> b -> Subst b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
t
  fmap a -> b
_ Subst a
PrimSubst = Subst b
forall t. Subst t
PrimSubst
  fmap a -> b
_ (SizeSubst DimDecl VName
v) = DimDecl VName -> Subst b
forall t. DimDecl VName -> Subst t
SizeSubst DimDecl VName
v

-- | Class of types which allow for substitution of types with no
-- annotations for type variable names.
class Substitutable a where
  applySubst :: TypeSubs -> a -> a

instance Substitutable (RetTypeBase (DimDecl VName) ()) where
  applySubst :: TypeSubs -> StructRetType -> StructRetType
applySubst TypeSubs
f (RetType [VName]
dims TypeBase (DimDecl VName) ()
t) =
    let RetType [VName]
more_dims TypeBase (DimDecl VName) ()
t' = TypeSubs -> TypeBase (DimDecl VName) () -> StructRetType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
substTypesRet TypeSubs
f TypeBase (DimDecl VName) ()
t
     in [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
more_dims) TypeBase (DimDecl VName) ()
t'

instance Substitutable (RetTypeBase (DimDecl VName) Aliasing) where
  applySubst :: TypeSubs
-> RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing
applySubst TypeSubs
f (RetType [VName]
dims PatType
t) =
    let RetType [VName]
more_dims PatType
t' = (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
substTypesRet VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
f' PatType
t
     in [VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
more_dims) PatType
t'
    where
      f' :: VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
f' = (Subst StructRetType
 -> Subst (RetTypeBase (DimDecl VName) Aliasing))
-> Maybe (Subst StructRetType)
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructRetType -> RetTypeBase (DimDecl VName) Aliasing)
-> Subst StructRetType
-> Subst (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((() -> Aliasing)
-> StructRetType -> RetTypeBase (DimDecl VName) Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliasing -> () -> Aliasing
forall a b. a -> b -> a
const Aliasing
forall a. Monoid a => a
mempty))) (Maybe (Subst StructRetType)
 -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> TypeSubs
-> VName
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
f

instance Substitutable (TypeBase (DimDecl VName) ()) where
  applySubst :: TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
applySubst = TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny

instance Substitutable (TypeBase (DimDecl VName) Aliasing) where
  applySubst :: TypeSubs -> PatType -> PatType
applySubst = (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> PatType -> PatType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny ((VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
 -> PatType -> PatType)
-> (TypeSubs
    -> VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> TypeSubs
-> PatType
-> PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Subst StructRetType
 -> Subst (RetTypeBase (DimDecl VName) Aliasing))
-> Maybe (Subst StructRetType)
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructRetType -> RetTypeBase (DimDecl VName) Aliasing)
-> Subst StructRetType
-> Subst (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((() -> Aliasing)
-> StructRetType -> RetTypeBase (DimDecl VName) Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliasing -> () -> Aliasing
forall a b. a -> b -> a
const Aliasing
forall a. Monoid a => a
mempty))) (Maybe (Subst StructRetType)
 -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> TypeSubs
-> VName
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.)

instance Substitutable (DimDecl VName) where
  applySubst :: TypeSubs -> DimDecl VName -> DimDecl VName
applySubst TypeSubs
f (NamedDim (QualName [VName]
_ VName
v))
    | Just (SizeSubst DimDecl VName
d) <- TypeSubs
f VName
v = DimDecl VName
d
  applySubst TypeSubs
_ DimDecl VName
d = DimDecl VName
d

instance Substitutable d => Substitutable (ShapeDecl d) where
  applySubst :: TypeSubs -> ShapeDecl d -> ShapeDecl d
applySubst TypeSubs
f = (d -> d) -> ShapeDecl d -> ShapeDecl d
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((d -> d) -> ShapeDecl d -> ShapeDecl d)
-> (d -> d) -> ShapeDecl d -> ShapeDecl d
forall a b. (a -> b) -> a -> b
$ TypeSubs -> d -> d
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
f

instance Substitutable Pat where
  applySubst :: TypeSubs -> Pat -> Pat
applySubst TypeSubs
f = Identity Pat -> Pat
forall a. Identity a -> a
runIdentity (Identity Pat -> Pat) -> (Pat -> Identity Pat) -> Pat -> Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASTMapper Identity -> Pat -> Identity Pat
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper Identity
mapper
    where
      mapper :: ASTMapper Identity
mapper =
        ASTMapper :: forall (m :: * -> *).
(ExpBase Info VName -> m (ExpBase Info VName))
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatType -> m PatType)
-> (StructRetType -> m StructRetType)
-> (RetTypeBase (DimDecl VName) Aliasing
    -> m (RetTypeBase (DimDecl VName) Aliasing))
-> ASTMapper m
ASTMapper
          { mapOnExp :: ExpBase Info VName -> Identity (ExpBase Info VName)
mapOnExp = ExpBase Info VName -> Identity (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return,
            mapOnName :: VName -> Identity VName
mapOnName = VName -> Identity VName
forall (m :: * -> *) a. Monad m => a -> m a
return,
            mapOnQualName :: QualName VName -> Identity (QualName VName)
mapOnQualName = QualName VName -> Identity (QualName VName)
forall (m :: * -> *) a. Monad m => a -> m a
return,
            mapOnStructType :: TypeBase (DimDecl VName) ()
-> Identity (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) ()
-> Identity (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) ()
 -> Identity (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> Identity (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
f,
            mapOnPatType :: PatType -> Identity PatType
mapOnPatType = PatType -> Identity PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> Identity PatType)
-> (PatType -> PatType) -> PatType -> Identity PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
f,
            mapOnStructRetType :: StructRetType -> Identity StructRetType
mapOnStructRetType = StructRetType -> Identity StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> Identity StructRetType)
-> (StructRetType -> StructRetType)
-> StructRetType
-> Identity StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
f,
            mapOnPatRetType :: RetTypeBase (DimDecl VName) Aliasing
-> Identity (RetTypeBase (DimDecl VName) Aliasing)
mapOnPatRetType = RetTypeBase (DimDecl VName) Aliasing
-> Identity (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) Aliasing
 -> Identity (RetTypeBase (DimDecl VName) Aliasing))
-> (RetTypeBase (DimDecl VName) Aliasing
    -> RetTypeBase (DimDecl VName) Aliasing)
-> RetTypeBase (DimDecl VName) Aliasing
-> Identity (RetTypeBase (DimDecl VName) Aliasing)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
f
          }

applyType ::
  Monoid als =>
  [TypeParam] ->
  TypeBase (DimDecl VName) als ->
  [StructTypeArg] ->
  TypeBase (DimDecl VName) als
applyType :: [TypeParam]
-> TypeBase (DimDecl VName) als
-> [StructTypeArg]
-> TypeBase (DimDecl VName) als
applyType [TypeParam]
ps TypeBase (DimDecl VName) als
t [StructTypeArg]
args = (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) als)))
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst (RetTypeBase (DimDecl VName) als))
-> Maybe (Subst (RetTypeBase (DimDecl VName) als))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (RetTypeBase (DimDecl VName) als))
substs) TypeBase (DimDecl VName) als
t
  where
    substs :: Map VName (Subst (RetTypeBase (DimDecl VName) als))
substs = [(VName, Subst (RetTypeBase (DimDecl VName) als))]
-> Map VName (Subst (RetTypeBase (DimDecl VName) als))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Subst (RetTypeBase (DimDecl VName) als))]
 -> Map VName (Subst (RetTypeBase (DimDecl VName) als)))
-> [(VName, Subst (RetTypeBase (DimDecl VName) als))]
-> Map VName (Subst (RetTypeBase (DimDecl VName) als))
forall a b. (a -> b) -> a -> b
$ (TypeParam
 -> StructTypeArg
 -> (VName, Subst (RetTypeBase (DimDecl VName) als)))
-> [TypeParam]
-> [StructTypeArg]
-> [(VName, Subst (RetTypeBase (DimDecl VName) als))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeParam
-> StructTypeArg
-> (VName, Subst (RetTypeBase (DimDecl VName) als))
forall as a.
(Monoid as, Eq a, IsName a) =>
TypeParamBase a
-> StructTypeArg -> (a, Subst (RetTypeBase (DimDecl VName) as))
mkSubst [TypeParam]
ps [StructTypeArg]
args
    -- We are assuming everything has already been type-checked for correctness.
    mkSubst :: TypeParamBase a
-> StructTypeArg -> (a, Subst (RetTypeBase (DimDecl VName) as))
mkSubst (TypeParamDim a
pv SrcLoc
_) (TypeArgDim DimDecl VName
d SrcLoc
_) =
      (a
pv, DimDecl VName -> Subst (RetTypeBase (DimDecl VName) as)
forall t. DimDecl VName -> Subst t
SizeSubst DimDecl VName
d)
    mkSubst (TypeParamType Liftedness
_ a
pv SrcLoc
_) (TypeArgType TypeBase (DimDecl VName) ()
at SrcLoc
_) =
      (a
pv, [TypeParam]
-> RetTypeBase (DimDecl VName) as
-> Subst (RetTypeBase (DimDecl VName) as)
forall t. [TypeParam] -> t -> Subst t
Subst [] (RetTypeBase (DimDecl VName) as
 -> Subst (RetTypeBase (DimDecl VName) as))
-> RetTypeBase (DimDecl VName) as
-> Subst (RetTypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as)
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ (() -> as)
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) as
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second () -> as
forall a. Monoid a => a
mempty TypeBase (DimDecl VName) ()
at)
    mkSubst TypeParamBase a
p StructTypeArg
a =
      String -> (a, Subst (RetTypeBase (DimDecl VName) as))
forall a. HasCallStack => String -> a
error (String -> (a, Subst (RetTypeBase (DimDecl VName) as)))
-> String -> (a, Subst (RetTypeBase (DimDecl VName) as))
forall a b. (a -> b) -> a -> b
$ String
"applyType mkSubst: cannot substitute " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StructTypeArg -> String
forall a. Pretty a => a -> String
pretty StructTypeArg
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" for " String -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeParamBase a -> String
forall a. Pretty a => a -> String
pretty TypeParamBase a
p

substTypesRet ::
  forall as.
  Monoid as =>
  (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))) ->
  TypeBase (DimDecl VName) as ->
  RetTypeBase (DimDecl VName) as
substTypesRet :: (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
substTypesRet VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))
lookupSubst TypeBase (DimDecl VName) as
ot =
  (TypeBase (DimDecl VName) as
 -> [VName] -> RetTypeBase (DimDecl VName) as)
-> (TypeBase (DimDecl VName) as, [VName])
-> RetTypeBase (DimDecl VName) as
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (([VName]
 -> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as)
-> TypeBase (DimDecl VName) as
-> [VName]
-> RetTypeBase (DimDecl VName) as
forall a b c. (a -> b -> c) -> b -> a -> c
flip [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType) ((TypeBase (DimDecl VName) as, [VName])
 -> RetTypeBase (DimDecl VName) as)
-> (TypeBase (DimDecl VName) as, [VName])
-> RetTypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ State [VName] (TypeBase (DimDecl VName) as)
-> [VName] -> (TypeBase (DimDecl VName) as, [VName])
forall s a. State s a -> s -> (a, s)
runState (TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType TypeBase (DimDecl VName) as
ot) []
  where
    -- In case we are substituting the same RetType in multiple
    -- places, we must ensure each instance is given distinct
    -- dimensions.  E.g. substituting 'a ↦ ?[n].[n]bool' into '(a,a)'
    -- should give '?[n][m].([n]bool,[m]bool)'.
    --
    -- XXX: the size names we invent here not globally unique.  This
    -- is _probably_ not a problem, since substituting types with
    -- outermost non-null existential sizes is done only when type
    -- checking modules.
    freshDims :: RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
freshDims (RetType [] TypeBase (DimDecl VName) as
t) = RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) as
 -> f (RetTypeBase (DimDecl VName) as))
-> RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase (DimDecl VName) as
t
    freshDims (RetType [VName]
ext TypeBase (DimDecl VName) as
t) = do
      [VName]
seen_ext <- f [VName]
forall s (m :: * -> *). MonadState s m => m s
get
      if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
seen_ext) [VName]
ext
        then RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) as
 -> f (RetTypeBase (DimDecl VName) as))
-> RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext TypeBase (DimDecl VName) as
t
        else do
          let start :: Int
start = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (VName -> Int) -> [VName] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Int
baseTag [VName]
seen_ext
              ext' :: [VName]
ext' = (Name -> Int -> VName) -> [Name] -> [Int] -> [VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Int -> VName
VName ((VName -> Name) -> [VName] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Name
baseName [VName]
ext) [Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 ..]
              extsubsts :: Map VName (Subst t)
extsubsts = [(VName, Subst t)] -> Map VName (Subst t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Subst t)] -> Map VName (Subst t))
-> [(VName, Subst t)] -> Map VName (Subst t)
forall a b. (a -> b) -> a -> b
$ [VName] -> [Subst t] -> [(VName, Subst t)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
ext ([Subst t] -> [(VName, Subst t)])
-> [Subst t] -> [(VName, Subst t)]
forall a b. (a -> b) -> a -> b
$ (VName -> Subst t) -> [VName] -> [Subst t]
forall a b. (a -> b) -> [a] -> [b]
map (DimDecl VName -> Subst t
forall t. DimDecl VName -> Subst t
SizeSubst (DimDecl VName -> Subst t)
-> (VName -> DimDecl VName) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> (VName -> QualName VName) -> VName -> DimDecl VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName) [VName]
ext'
              RetType [] TypeBase (DimDecl VName) as
t' = (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
substTypesRet (VName
-> Map VName (Subst (RetTypeBase (DimDecl VName) as))
-> Maybe (Subst (RetTypeBase (DimDecl VName) as))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (RetTypeBase (DimDecl VName) as))
forall t. Map VName (Subst t)
extsubsts) TypeBase (DimDecl VName) as
t
          RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) as
 -> f (RetTypeBase (DimDecl VName) as))
-> RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext' TypeBase (DimDecl VName) as
t'

    onType :: TypeBase (DimDecl VName) as -> State [VName] (TypeBase (DimDecl VName) as)

    onType :: TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType (Array as
als Uniqueness
u ScalarTypeBase (DimDecl VName) ()
et ShapeDecl (DimDecl VName)
shape) = do
      TypeBase (DimDecl VName) as
t <-
        TypeBase (DimDecl VName) as
-> ShapeDecl (DimDecl VName)
-> Uniqueness
-> TypeBase (DimDecl VName) as
forall as dim.
Monoid as =>
TypeBase dim as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOf (TypeBase (DimDecl VName) as
 -> ShapeDecl (DimDecl VName)
 -> Uniqueness
 -> TypeBase (DimDecl VName) as)
-> State [VName] (TypeBase (DimDecl VName) as)
-> StateT
     [VName]
     Identity
     (ShapeDecl (DimDecl VName)
      -> Uniqueness -> TypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
et TypeBase (DimDecl VName) () -> as -> TypeBase (DimDecl VName) as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` as
forall a. Monoid a => a
mempty)
          StateT
  [VName]
  Identity
  (ShapeDecl (DimDecl VName)
   -> Uniqueness -> TypeBase (DimDecl VName) as)
-> StateT [VName] Identity (ShapeDecl (DimDecl VName))
-> StateT
     [VName] Identity (Uniqueness -> TypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ShapeDecl (DimDecl VName)
-> StateT [VName] Identity (ShapeDecl (DimDecl VName))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeSubs -> ShapeDecl (DimDecl VName) -> ShapeDecl (DimDecl VName)
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
lookupSubst' ShapeDecl (DimDecl VName)
shape)
          StateT [VName] Identity (Uniqueness -> TypeBase (DimDecl VName) as)
-> StateT [VName] Identity Uniqueness
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Uniqueness -> StateT [VName] Identity Uniqueness
forall (f :: * -> *) a. Applicative f => a -> f a
pure Uniqueness
u
      TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) as
t TypeBase (DimDecl VName) as -> as -> TypeBase (DimDecl VName) as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` as
als
    onType (Scalar (Prim PrimType
t)) = TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) as
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    onType (Scalar (TypeVar as
als Uniqueness
u TypeName
v [StructTypeArg]
targs)) = do
      [StructTypeArg]
targs' <- (StructTypeArg -> StateT [VName] Identity StructTypeArg)
-> [StructTypeArg] -> StateT [VName] Identity [StructTypeArg]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM StructTypeArg -> StateT [VName] Identity StructTypeArg
forall (m :: * -> *).
MonadState [VName] m =>
StructTypeArg -> m StructTypeArg
subsTypeArg [StructTypeArg]
targs
      case VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))
lookupSubst (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf (TypeName -> QualName VName
qualNameFromTypeName TypeName
v) of
        Just (Subst [TypeParam]
ps RetTypeBase (DimDecl VName) as
rt) -> do
          RetType [VName]
ext TypeBase (DimDecl VName) as
t <- RetTypeBase (DimDecl VName) as
-> StateT [VName] Identity (RetTypeBase (DimDecl VName) as)
forall (f :: * -> *) as.
(MonadState [VName] f, Monoid as) =>
RetTypeBase (DimDecl VName) as
-> f (RetTypeBase (DimDecl VName) as)
freshDims RetTypeBase (DimDecl VName) as
rt
          ([VName] -> [VName]) -> StateT [VName] Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ([VName]
ext [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++)
          TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$
            [TypeParam]
-> TypeBase (DimDecl VName) as
-> [StructTypeArg]
-> TypeBase (DimDecl VName) as
forall als.
Monoid als =>
[TypeParam]
-> TypeBase (DimDecl VName) als
-> [StructTypeArg]
-> TypeBase (DimDecl VName) als
applyType [TypeParam]
ps (TypeBase (DimDecl VName) as
t TypeBase (DimDecl VName) as -> as -> TypeBase (DimDecl VName) as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` as
forall a. Monoid a => a
mempty) [StructTypeArg]
targs'
              TypeBase (DimDecl VName) as
-> Uniqueness -> TypeBase (DimDecl VName) as
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u TypeBase (DimDecl VName) as
-> (as -> as) -> TypeBase (DimDecl VName) as
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` (as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als)
        Just Subst (RetTypeBase (DimDecl VName) as)
PrimSubst ->
          TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ as
-> Uniqueness
-> TypeName
-> [StructTypeArg]
-> ScalarTypeBase (DimDecl VName) as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar as
forall a. Monoid a => a
mempty Uniqueness
u TypeName
v [StructTypeArg]
targs'
        Maybe (Subst (RetTypeBase (DimDecl VName) as))
_ ->
          TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall a b. (a -> b) -> a -> b
$ as
-> Uniqueness
-> TypeName
-> [StructTypeArg]
-> ScalarTypeBase (DimDecl VName) as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar as
als Uniqueness
u TypeName
v [StructTypeArg]
targs'
    onType (Scalar (Record Map Name (TypeBase (DimDecl VName) as)
ts)) =
      ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> (Map Name (TypeBase (DimDecl VName) as)
    -> ScalarTypeBase (DimDecl VName) as)
-> Map Name (TypeBase (DimDecl VName) as)
-> TypeBase (DimDecl VName) as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase (DimDecl VName) as)
-> ScalarTypeBase (DimDecl VName) as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase (DimDecl VName) as)
 -> TypeBase (DimDecl VName) as)
-> StateT [VName] Identity (Map Name (TypeBase (DimDecl VName) as))
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> Map Name (TypeBase (DimDecl VName) as)
-> StateT [VName] Identity (Map Name (TypeBase (DimDecl VName) as))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType Map Name (TypeBase (DimDecl VName) as)
ts
    onType (Scalar (Arrow as
als PName
v TypeBase (DimDecl VName) as
t1 RetTypeBase (DimDecl VName) as
t2)) =
      ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> StateT [VName] Identity (ScalarTypeBase (DimDecl VName) as)
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (as
-> PName
-> TypeBase (DimDecl VName) as
-> RetTypeBase (DimDecl VName) as
-> ScalarTypeBase (DimDecl VName) as
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
als PName
v (TypeBase (DimDecl VName) as
 -> RetTypeBase (DimDecl VName) as
 -> ScalarTypeBase (DimDecl VName) as)
-> State [VName] (TypeBase (DimDecl VName) as)
-> StateT
     [VName]
     Identity
     (RetTypeBase (DimDecl VName) as
      -> ScalarTypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType TypeBase (DimDecl VName) as
t1 StateT
  [VName]
  Identity
  (RetTypeBase (DimDecl VName) as
   -> ScalarTypeBase (DimDecl VName) as)
-> StateT [VName] Identity (RetTypeBase (DimDecl VName) as)
-> StateT [VName] Identity (ScalarTypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RetTypeBase (DimDecl VName) as
-> StateT [VName] Identity (RetTypeBase (DimDecl VName) as)
forall (m :: * -> *).
MonadState [VName] m =>
RetTypeBase (DimDecl VName) as
-> m (RetTypeBase (DimDecl VName) as)
onRetType RetTypeBase (DimDecl VName) as
t2)
    onType (Scalar (Sum Map Name [TypeBase (DimDecl VName) as]
ts)) =
      ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as)
-> (Map Name [TypeBase (DimDecl VName) as]
    -> ScalarTypeBase (DimDecl VName) as)
-> Map Name [TypeBase (DimDecl VName) as]
-> TypeBase (DimDecl VName) as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase (DimDecl VName) as]
-> ScalarTypeBase (DimDecl VName) as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase (DimDecl VName) as]
 -> TypeBase (DimDecl VName) as)
-> StateT [VName] Identity (Map Name [TypeBase (DimDecl VName) as])
-> State [VName] (TypeBase (DimDecl VName) as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase (DimDecl VName) as]
 -> StateT [VName] Identity [TypeBase (DimDecl VName) as])
-> Map Name [TypeBase (DimDecl VName) as]
-> StateT [VName] Identity (Map Name [TypeBase (DimDecl VName) as])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase (DimDecl VName) as
 -> State [VName] (TypeBase (DimDecl VName) as))
-> [TypeBase (DimDecl VName) as]
-> StateT [VName] Identity [TypeBase (DimDecl VName) as]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType) Map Name [TypeBase (DimDecl VName) as]
ts

    onRetType :: RetTypeBase (DimDecl VName) as
-> m (RetTypeBase (DimDecl VName) as)
onRetType (RetType [VName]
dims TypeBase (DimDecl VName) as
t) = do
      [VName]
ext <- m [VName]
forall s (m :: * -> *). MonadState s m => m s
get
      let (TypeBase (DimDecl VName) as
t', [VName]
ext') = State [VName] (TypeBase (DimDecl VName) as)
-> [VName] -> (TypeBase (DimDecl VName) as, [VName])
forall s a. State s a -> s -> (a, s)
runState (TypeBase (DimDecl VName) as
-> State [VName] (TypeBase (DimDecl VName) as)
onType TypeBase (DimDecl VName) as
t) [VName]
ext
          new_ext :: [VName]
new_ext = [VName]
ext' [VName] -> [VName] -> [VName]
forall a. Eq a => [a] -> [a] -> [a]
\\ [VName]
ext
      case TypeBase (DimDecl VName) as
t of
        Scalar Arrow {} -> do
          [VName] -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put [VName]
ext'
          RetTypeBase (DimDecl VName) as
-> m (RetTypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) as
 -> m (RetTypeBase (DimDecl VName) as))
-> RetTypeBase (DimDecl VName) as
-> m (RetTypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase (DimDecl VName) as
t'
        TypeBase (DimDecl VName) as
_ ->
          RetTypeBase (DimDecl VName) as
-> m (RetTypeBase (DimDecl VName) as)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) as
 -> m (RetTypeBase (DimDecl VName) as))
-> RetTypeBase (DimDecl VName) as
-> m (RetTypeBase (DimDecl VName) as)
forall a b. (a -> b) -> a -> b
$ [VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
new_ext [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims) TypeBase (DimDecl VName) as
t'

    subsTypeArg :: StructTypeArg -> m StructTypeArg
subsTypeArg (TypeArgType TypeBase (DimDecl VName) ()
t SrcLoc
loc) = do
      let RetType [VName]
dims TypeBase (DimDecl VName) ()
t' = TypeSubs -> TypeBase (DimDecl VName) () -> StructRetType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
substTypesRet TypeSubs
lookupSubst' TypeBase (DimDecl VName) ()
t
      ([VName] -> [VName]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ([VName]
dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++)
      StructTypeArg -> m StructTypeArg
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructTypeArg -> m StructTypeArg)
-> StructTypeArg -> m StructTypeArg
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> SrcLoc -> StructTypeArg
forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType TypeBase (DimDecl VName) ()
t' SrcLoc
loc
    subsTypeArg (TypeArgDim DimDecl VName
v SrcLoc
loc) =
      StructTypeArg -> m StructTypeArg
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructTypeArg -> m StructTypeArg)
-> StructTypeArg -> m StructTypeArg
forall a b. (a -> b) -> a -> b
$ DimDecl VName -> SrcLoc -> StructTypeArg
forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim (TypeSubs -> DimDecl VName -> DimDecl VName
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
lookupSubst' DimDecl VName
v) SrcLoc
loc

    lookupSubst' :: TypeSubs
lookupSubst' = (Subst (RetTypeBase (DimDecl VName) as) -> Subst StructRetType)
-> Maybe (Subst (RetTypeBase (DimDecl VName) as))
-> Maybe (Subst StructRetType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((RetTypeBase (DimDecl VName) as -> StructRetType)
-> Subst (RetTypeBase (DimDecl VName) as) -> Subst StructRetType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((RetTypeBase (DimDecl VName) as -> StructRetType)
 -> Subst (RetTypeBase (DimDecl VName) as) -> Subst StructRetType)
-> (RetTypeBase (DimDecl VName) as -> StructRetType)
-> Subst (RetTypeBase (DimDecl VName) as)
-> Subst StructRetType
forall a b. (a -> b) -> a -> b
$ (as -> ()) -> RetTypeBase (DimDecl VName) as -> StructRetType
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (() -> as -> ()
forall a b. a -> b -> a
const ())) (Maybe (Subst (RetTypeBase (DimDecl VName) as))
 -> Maybe (Subst StructRetType))
-> (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeSubs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))
lookupSubst

-- | Perform substitutions, from type names to types, on a type. Works
-- regardless of what shape and uniqueness information is attached to the type.
substTypesAny ::
  Monoid as =>
  (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))) ->
  TypeBase (DimDecl VName) as ->
  TypeBase (DimDecl VName) as
substTypesAny :: (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))
lookupSubst TypeBase (DimDecl VName) as
ot =
  case (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
substTypesRet VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as))
lookupSubst TypeBase (DimDecl VName) as
ot of
    RetType [] TypeBase (DimDecl VName) as
ot' -> TypeBase (DimDecl VName) as
ot'
    RetType [VName]
dims TypeBase (DimDecl VName) as
ot' ->
      -- XXX HACK FIXME: turn any sizes that propagate to the top into
      -- AnyDim.  This should _never_ happen during type-checking, but
      -- may happen as we substitute types during monomorphisation and
      -- defunctorisation later on. See Note [AnyDim]
      let toAny :: DimDecl VName -> DimDecl VName
toAny (NamedDim QualName VName
v)
            | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
dims = Maybe VName -> DimDecl VName
forall vn. Maybe vn -> DimDecl vn
AnyDim Maybe VName
forall a. Maybe a
Nothing
          toAny DimDecl VName
d = DimDecl VName
d
       in (DimDecl VName -> DimDecl VName)
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
toAny TypeBase (DimDecl VName) as
ot'

-- Note [AnyDim]
--
-- Consider a program:
--
-- module m : { type~ t } = { type~ t = ?[n].[n]bool }
-- let f (x: m.t) (y: m.t) = 0
--
-- After defunctorisation (and inlining the definitions of types), we
-- want this:
--
-- let f [n][m] (x: [n]bool) (y: [m]bool) = 0
--
-- But this means that defunctorisation would need to redo some amount
-- of size inference.  Not so complicated in the example above, but
-- what if loops and branches are involved?
--
-- So instead, what defunctorisation actually does is produce this:
--
-- let f (x: []bool) (y: []bool) = 0
--
-- I.e. we put in empty dimensions (AnyDim), which are much later
-- turned into distinct sizes in Futhark.Internalise.Exps.  This will
-- result in unnecessary dynamic size checks, which will hopefully be
-- optimised away.
--
-- Important: The type checker will _never_ produce programs with
-- AnyDim, but unfortunately some of the compilation steps
-- (defunctorisation, monomorphisation, defunctionalisation) will do
-- so.  Similarly, the core language is also perfectly well behaved.
--
-- Example with monomorphisation:
--
-- let f '~a (b: bool) (x: () -> a) (y: () -> a) : a = if b then x () else y ()
-- let g = f true (\() -> [1]) (\() -> [1,2])
--
-- This should produce:
--
-- let f (b: bool) (x: () -> ?[n].[n]i32) (y: () -> ?[m].[m]i32) : ?[k].[k]i32 =
--   if b then x () else y ()
--
-- Not so easy!  Again, what we actually produce is
--
-- let f (b: bool) (x: () -> []i32) (y: () -> []i32) : []i32 =
--   if b then x () else y ()