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

-- | This module provides various simple ways to query and manipulate
-- fundamental Futhark terms, such as types and values.  The intent is to
-- keep "Futhark.Language.Syntax" simple, and put whatever embellishments
-- we need here.
module Language.Futhark.Prop
  ( -- * Various
    Intrinsic (..),
    intrinsics,
    maxIntrinsicTag,
    namesToPrimTypes,
    qualName,
    qualify,
    typeName,
    valueType,
    primValueType,
    leadingOperator,
    progImports,
    decImports,
    progModuleTypes,
    identifierReference,
    prettyStacktrace,

    -- * Queries on expressions
    typeOf,

    -- * Queries on patterns and params
    patternIdents,
    patternNames,
    patternMap,
    patternType,
    patternStructType,
    patternParam,
    patternOrderZero,
    patternDimNames,

    -- * Queries on types
    uniqueness,
    unique,
    aliases,
    diet,
    arrayRank,
    arrayShape,
    nestedDims,
    orderZero,
    unfoldFunType,
    foldFunType,
    typeVars,
    typeDimNames,
    primByteSize,

    -- * Operations on types
    rank,
    peelArray,
    stripArray,
    arrayOf,
    toStructural,
    toStruct,
    fromStruct,
    setAliases,
    addAliases,
    setUniqueness,
    noSizes,
    addSizes,
    anySizes,
    traverseDims,
    DimPos (..),
    mustBeExplicit,
    mustBeExplicitInType,
    tupleRecord,
    isTupleRecord,
    areTupleFields,
    tupleFields,
    tupleFieldNames,
    sortFields,
    sortConstrs,
    isTypeParam,
    isSizeParam,
    combineTypeShapes,
    matchDims,
    unscopeType,
    onRecordField,
    -- | Values of these types are produces by the parser.  They use
    -- unadorned names and have no type information, apart from that
    -- which is syntactically required.
    NoInfo (..),
    UncheckedType,
    UncheckedTypeExp,
    UncheckedIdent,
    UncheckedTypeDecl,
    UncheckedDimIndex,
    UncheckedExp,
    UncheckedModExp,
    UncheckedSigExp,
    UncheckedTypeParam,
    UncheckedPattern,
    UncheckedValBind,
    UncheckedDec,
    UncheckedSpec,
    UncheckedProg,
    UncheckedCase,
  )
where

import Control.Monad.State
import Data.Bifoldable
import Data.Bifunctor
import Data.Bitraversable (bitraverse)
import Data.Char
import Data.Foldable
import Data.List (genericLength, isPrefixOf, nub, sortOn)
import qualified Data.Map.Strict as M
import Data.Maybe
import Data.Ord
import qualified Data.Set as S
import qualified Futhark.IR.Primitive as Primitive
import Futhark.Util (maxinum)
import Futhark.Util.Pretty
import Language.Futhark.Syntax

-- | Return the dimensionality of a type.  For non-arrays, this is
-- zero.  For a one-dimensional array it is one, for a two-dimensional
-- it is two, and so forth.
arrayRank :: TypeBase dim as -> Int
arrayRank :: TypeBase dim as -> Int
arrayRank = ShapeDecl dim -> Int
forall dim. ShapeDecl dim -> Int
shapeRank (ShapeDecl dim -> Int)
-> (TypeBase dim as -> ShapeDecl dim) -> TypeBase dim as -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase dim as -> ShapeDecl dim
forall dim as. TypeBase dim as -> ShapeDecl dim
arrayShape

-- | Return the shape of a type - for non-arrays, this is 'mempty'.
arrayShape :: TypeBase dim as -> ShapeDecl dim
arrayShape :: TypeBase dim as -> ShapeDecl dim
arrayShape (Array as
_ Uniqueness
_ ScalarTypeBase dim ()
_ ShapeDecl dim
ds) = ShapeDecl dim
ds
arrayShape TypeBase dim as
_ = ShapeDecl dim
forall a. Monoid a => a
mempty

-- | Return any shape declarations in the type, with duplicates
-- removed.
nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims TypeBase (DimDecl VName) as
t =
  case TypeBase (DimDecl VName) as
t of
    Array as
_ Uniqueness
_ ScalarTypeBase (DimDecl VName) ()
a ShapeDecl (DimDecl VName)
ds ->
      [DimDecl VName] -> [DimDecl VName]
forall a. Eq a => [a] -> [a]
nub ([DimDecl VName] -> [DimDecl VName])
-> [DimDecl VName] -> [DimDecl VName]
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
a) [DimDecl VName] -> [DimDecl VName] -> [DimDecl VName]
forall a. Semigroup a => a -> a -> a
<> ShapeDecl (DimDecl VName) -> [DimDecl VName]
forall dim. ShapeDecl dim -> [dim]
shapeDims ShapeDecl (DimDecl VName)
ds
    Scalar (Record Map Name (TypeBase (DimDecl VName) as)
fs) ->
      [DimDecl VName] -> [DimDecl VName]
forall a. Eq a => [a] -> [a]
nub ([DimDecl VName] -> [DimDecl VName])
-> [DimDecl VName] -> [DimDecl VName]
forall a b. (a -> b) -> a -> b
$ (TypeBase (DimDecl VName) as -> [DimDecl VName])
-> Map Name (TypeBase (DimDecl VName) as) -> [DimDecl VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase (DimDecl VName) as -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims Map Name (TypeBase (DimDecl VName) as)
fs
    Scalar Prim {} ->
      [DimDecl VName]
forall a. Monoid a => a
mempty
    Scalar (Sum Map Name [TypeBase (DimDecl VName) as]
cs) ->
      [DimDecl VName] -> [DimDecl VName]
forall a. Eq a => [a] -> [a]
nub ([DimDecl VName] -> [DimDecl VName])
-> [DimDecl VName] -> [DimDecl VName]
forall a b. (a -> b) -> a -> b
$ ([TypeBase (DimDecl VName) as] -> [DimDecl VName])
-> Map Name [TypeBase (DimDecl VName) as] -> [DimDecl VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((TypeBase (DimDecl VName) as -> [DimDecl VName])
-> [TypeBase (DimDecl VName) as] -> [DimDecl VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase (DimDecl VName) as -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims) Map Name [TypeBase (DimDecl VName) as]
cs
    Scalar (Arrow as
_ PName
v TypeBase (DimDecl VName) as
t1 TypeBase (DimDecl VName) as
t2) ->
      (DimDecl VName -> Bool) -> [DimDecl VName] -> [DimDecl VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (PName -> DimDecl VName -> Bool
notV PName
v) ([DimDecl VName] -> [DimDecl VName])
-> [DimDecl VName] -> [DimDecl VName]
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) as -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims TypeBase (DimDecl VName) as
t1 [DimDecl VName] -> [DimDecl VName] -> [DimDecl VName]
forall a. Semigroup a => a -> a -> a
<> TypeBase (DimDecl VName) as -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims TypeBase (DimDecl VName) as
t2
    Scalar (TypeVar as
_ Uniqueness
_ TypeName
_ [TypeArg (DimDecl VName)]
targs) ->
      (TypeArg (DimDecl VName) -> [DimDecl VName])
-> [TypeArg (DimDecl VName)] -> [DimDecl VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap TypeArg (DimDecl VName) -> [DimDecl VName]
typeArgDims [TypeArg (DimDecl VName)]
targs
  where
    typeArgDims :: TypeArg (DimDecl VName) -> [DimDecl VName]
typeArgDims (TypeArgDim DimDecl VName
d SrcLoc
_) = [DimDecl VName
d]
    typeArgDims (TypeArgType TypeBase (DimDecl VName) ()
at SrcLoc
_) = TypeBase (DimDecl VName) () -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims TypeBase (DimDecl VName) ()
at

    notV :: PName -> DimDecl VName -> Bool
notV PName
Unnamed = Bool -> DimDecl VName -> Bool
forall a b. a -> b -> a
const Bool
True
    notV (Named VName
v) = (DimDecl VName -> DimDecl VName -> Bool
forall a. Eq a => a -> a -> Bool
/= QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v))

-- | Change the shape of a type to be just the rank.
noSizes :: TypeBase (DimDecl vn) as -> TypeBase () as
noSizes :: TypeBase (DimDecl vn) as -> TypeBase () as
noSizes = (DimDecl vn -> ()) -> TypeBase (DimDecl vn) as -> TypeBase () as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((DimDecl vn -> ()) -> TypeBase (DimDecl vn) as -> TypeBase () as)
-> (DimDecl vn -> ()) -> TypeBase (DimDecl vn) as -> TypeBase () as
forall a b. (a -> b) -> a -> b
$ () -> DimDecl vn -> ()
forall a b. a -> b -> a
const ()

-- | Add size annotations that are all 'AnyDim'.
addSizes :: TypeBase () as -> TypeBase (DimDecl vn) as
addSizes :: TypeBase () as -> TypeBase (DimDecl vn) as
addSizes = (() -> DimDecl vn) -> TypeBase () as -> TypeBase (DimDecl vn) as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((() -> DimDecl vn) -> TypeBase () as -> TypeBase (DimDecl vn) as)
-> (() -> DimDecl vn) -> TypeBase () as -> TypeBase (DimDecl vn) as
forall a b. (a -> b) -> a -> b
$ DimDecl vn -> () -> DimDecl vn
forall a b. a -> b -> a
const DimDecl vn
forall vn. DimDecl vn
AnyDim

-- | Change all size annotations to be 'AnyDim'.
anySizes :: TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as
anySizes :: TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as
anySizes = (DimDecl vn -> DimDecl vn)
-> TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((DimDecl vn -> DimDecl vn)
 -> TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as)
-> (DimDecl vn -> DimDecl vn)
-> TypeBase (DimDecl vn) as
-> TypeBase (DimDecl vn) as
forall a b. (a -> b) -> a -> b
$ DimDecl vn -> DimDecl vn -> DimDecl vn
forall a b. a -> b -> a
const DimDecl vn
forall vn. DimDecl vn
AnyDim

-- | Where does this dimension occur?
data DimPos
  = -- | Immediately in the argument to 'traverseDims'.
    PosImmediate
  | -- | In a function parameter type.
    PosParam
  | -- | In a function return type.
    PosReturn
  deriving (DimPos -> DimPos -> Bool
(DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool) -> Eq DimPos
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DimPos -> DimPos -> Bool
$c/= :: DimPos -> DimPos -> Bool
== :: DimPos -> DimPos -> Bool
$c== :: DimPos -> DimPos -> Bool
Eq, Eq DimPos
Eq DimPos
-> (DimPos -> DimPos -> Ordering)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> DimPos)
-> (DimPos -> DimPos -> DimPos)
-> Ord DimPos
DimPos -> DimPos -> Bool
DimPos -> DimPos -> Ordering
DimPos -> DimPos -> DimPos
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DimPos -> DimPos -> DimPos
$cmin :: DimPos -> DimPos -> DimPos
max :: DimPos -> DimPos -> DimPos
$cmax :: DimPos -> DimPos -> DimPos
>= :: DimPos -> DimPos -> Bool
$c>= :: DimPos -> DimPos -> Bool
> :: DimPos -> DimPos -> Bool
$c> :: DimPos -> DimPos -> Bool
<= :: DimPos -> DimPos -> Bool
$c<= :: DimPos -> DimPos -> Bool
< :: DimPos -> DimPos -> Bool
$c< :: DimPos -> DimPos -> Bool
compare :: DimPos -> DimPos -> Ordering
$ccompare :: DimPos -> DimPos -> Ordering
$cp1Ord :: Eq DimPos
Ord, Int -> DimPos -> ShowS
[DimPos] -> ShowS
DimPos -> String
(Int -> DimPos -> ShowS)
-> (DimPos -> String) -> ([DimPos] -> ShowS) -> Show DimPos
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DimPos] -> ShowS
$cshowList :: [DimPos] -> ShowS
show :: DimPos -> String
$cshow :: DimPos -> String
showsPrec :: Int -> DimPos -> ShowS
$cshowsPrec :: Int -> DimPos -> ShowS
Show)

-- | Perform a traversal (possibly including replacement) on sizes
-- that are parameters in a function type, but also including the type
-- immediately passed to the function.  Also passes along a set of the
-- parameter names inside the type that have come in scope at the
-- occurrence of the dimension.
traverseDims ::
  forall f fdim tdim als.
  Applicative f =>
  (S.Set VName -> DimPos -> fdim -> f tdim) ->
  TypeBase fdim als ->
  f (TypeBase tdim als)
traverseDims :: (Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName -> DimPos -> fdim -> f tdim
f = Set VName -> DimPos -> TypeBase fdim als -> f (TypeBase tdim als)
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
forall a. Monoid a => a
mempty DimPos
PosImmediate
  where
    go ::
      forall als'.
      S.Set VName ->
      DimPos ->
      TypeBase fdim als' ->
      f (TypeBase tdim als')
    go :: Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b t :: TypeBase fdim als'
t@Array {} =
      (fdim -> f tdim)
-> (als' -> f als') -> TypeBase fdim als' -> f (TypeBase tdim als')
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (Set VName -> DimPos -> fdim -> f tdim
f Set VName
bound DimPos
b) als' -> f als'
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase fdim als'
t
    go Set VName
bound DimPos
b (Scalar (Record Map Name (TypeBase fdim als')
fields)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> (Map Name (TypeBase tdim als') -> ScalarTypeBase tdim als')
-> Map Name (TypeBase tdim als')
-> TypeBase tdim als'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase tdim als') -> ScalarTypeBase tdim als'
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase tdim als') -> TypeBase tdim als')
-> f (Map Name (TypeBase tdim als')) -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase fdim als' -> f (TypeBase tdim als'))
-> Map Name (TypeBase fdim als')
-> f (Map Name (TypeBase tdim als'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b) Map Name (TypeBase fdim als')
fields
    go Set VName
bound DimPos
b (Scalar (TypeVar als'
as Uniqueness
u TypeName
tn [TypeArg fdim]
targs)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> f (ScalarTypeBase tdim als') -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (als'
-> Uniqueness
-> TypeName
-> [TypeArg tdim]
-> ScalarTypeBase tdim als'
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar als'
as Uniqueness
u TypeName
tn ([TypeArg tdim] -> ScalarTypeBase tdim als')
-> f [TypeArg tdim] -> f (ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg fdim -> f (TypeArg tdim))
-> [TypeArg fdim] -> f [TypeArg tdim]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Set VName -> DimPos -> TypeArg fdim -> f (TypeArg tdim)
onTypeArg Set VName
bound DimPos
b) [TypeArg fdim]
targs)
    go Set VName
bound DimPos
b (Scalar (Sum Map Name [TypeBase fdim als']
cs)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> (Map Name [TypeBase tdim als'] -> ScalarTypeBase tdim als')
-> Map Name [TypeBase tdim als']
-> TypeBase tdim als'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase tdim als'] -> ScalarTypeBase tdim als'
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase tdim als'] -> TypeBase tdim als')
-> f (Map Name [TypeBase tdim als']) -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase fdim als'] -> f [TypeBase tdim als'])
-> Map Name [TypeBase fdim als']
-> f (Map Name [TypeBase tdim als'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase fdim als' -> f (TypeBase tdim als'))
-> [TypeBase fdim als'] -> f [TypeBase tdim als']
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b)) Map Name [TypeBase fdim als']
cs
    go Set VName
_ DimPos
_ (Scalar (Prim PrimType
t)) =
      TypeBase tdim als' -> f (TypeBase tdim als')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase tdim als' -> f (TypeBase tdim als'))
-> TypeBase tdim als' -> f (TypeBase tdim als')
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> ScalarTypeBase tdim als' -> TypeBase tdim als'
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase tdim als'
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    go Set VName
bound DimPos
_ (Scalar (Arrow als'
als PName
p TypeBase fdim als'
t1 TypeBase fdim als'
t2)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> f (ScalarTypeBase tdim als') -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (als'
-> PName
-> TypeBase tdim als'
-> TypeBase tdim als'
-> ScalarTypeBase tdim als'
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow als'
als PName
p (TypeBase tdim als'
 -> TypeBase tdim als' -> ScalarTypeBase tdim als')
-> f (TypeBase tdim als')
-> f (TypeBase tdim als' -> ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosParam TypeBase fdim als'
t1 f (TypeBase tdim als' -> ScalarTypeBase tdim als')
-> f (TypeBase tdim als') -> f (ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosReturn TypeBase fdim als'
t2)
      where
        bound' :: Set VName
bound' = case PName
p of
          Named VName
p' -> VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert VName
p' Set VName
bound
          PName
Unnamed -> Set VName
bound

    onTypeArg :: Set VName -> DimPos -> TypeArg fdim -> f (TypeArg tdim)
onTypeArg Set VName
bound DimPos
b (TypeArgDim fdim
d SrcLoc
loc) =
      tdim -> SrcLoc -> TypeArg tdim
forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim (tdim -> SrcLoc -> TypeArg tdim)
-> f tdim -> f (SrcLoc -> TypeArg tdim)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> DimPos -> fdim -> f tdim
f Set VName
bound DimPos
b fdim
d f (SrcLoc -> TypeArg tdim) -> f SrcLoc -> f (TypeArg tdim)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> f SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    onTypeArg Set VName
bound DimPos
b (TypeArgType TypeBase fdim ()
t SrcLoc
loc) =
      TypeBase tdim () -> SrcLoc -> TypeArg tdim
forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (TypeBase tdim () -> SrcLoc -> TypeArg tdim)
-> f (TypeBase tdim ()) -> f (SrcLoc -> TypeArg tdim)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> DimPos -> TypeBase fdim () -> f (TypeBase tdim ())
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b TypeBase fdim ()
t f (SrcLoc -> TypeArg tdim) -> f SrcLoc -> f (TypeArg tdim)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> f SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

mustBeExplicitAux :: StructType -> M.Map VName Bool
mustBeExplicitAux :: TypeBase (DimDecl VName) () -> Map VName Bool
mustBeExplicitAux TypeBase (DimDecl VName) ()
t =
  State (Map VName Bool) (TypeBase () ())
-> Map VName Bool -> Map VName Bool
forall s a. State s a -> s -> s
execState ((Set VName
 -> DimPos -> DimDecl VName -> StateT (Map VName Bool) Identity ())
-> TypeBase (DimDecl VName) ()
-> State (Map VName Bool) (TypeBase () ())
forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName
-> DimPos -> DimDecl VName -> StateT (Map VName Bool) Identity ()
forall k (m :: * -> *).
(Ord k, MonadState (Map k Bool) m) =>
Set k -> DimPos -> DimDecl k -> m ()
onDim TypeBase (DimDecl VName) ()
t) Map VName Bool
forall a. Monoid a => a
mempty
  where
    onDim :: Set k -> DimPos -> DimDecl k -> m ()
onDim Set k
bound DimPos
_ (NamedDim QualName k
d)
      | QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
d k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set k
bound =
        (Map k Bool -> Map k Bool) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map k Bool -> Map k Bool) -> m ())
-> (Map k Bool -> Map k Bool) -> m ()
forall a b. (a -> b) -> a -> b
$ \Map k Bool
s -> (Bool -> Bool -> Bool) -> k -> Bool -> Map k Bool -> Map k Bool
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Bool -> Bool -> Bool
(&&) (QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
d) Bool
False Map k Bool
s
    onDim Set k
_ DimPos
PosImmediate (NamedDim QualName k
d) =
      (Map k Bool -> Map k Bool) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map k Bool -> Map k Bool) -> m ())
-> (Map k Bool -> Map k Bool) -> m ()
forall a b. (a -> b) -> a -> b
$ \Map k Bool
s -> (Bool -> Bool -> Bool) -> k -> Bool -> Map k Bool -> Map k Bool
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Bool -> Bool -> Bool
(&&) (QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
d) Bool
False Map k Bool
s
    onDim Set k
_ DimPos
_ (NamedDim QualName k
d) =
      (Map k Bool -> Map k Bool) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map k Bool -> Map k Bool) -> m ())
-> (Map k Bool -> Map k Bool) -> m ()
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool -> Bool) -> k -> Bool -> Map k Bool -> Map k Bool
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Bool -> Bool -> Bool
(&&) (QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
d) Bool
True
    onDim Set k
_ DimPos
_ DimDecl k
_ =
      () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Figure out which of the sizes in a parameter type must be passed
-- explicitly, because their first use is as something else than just
-- an array dimension.  'mustBeExplicit' is like this function, but
-- first decomposes into parameter types.
mustBeExplicitInType :: StructType -> S.Set VName
mustBeExplicitInType :: TypeBase (DimDecl VName) () -> Set VName
mustBeExplicitInType TypeBase (DimDecl VName) ()
t =
  [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ Map VName Bool -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName Bool -> [VName]) -> Map VName Bool -> [VName]
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> Map VName Bool -> Map VName Bool
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter Bool -> Bool
forall a. a -> a
id (Map VName Bool -> Map VName Bool)
-> Map VName Bool -> Map VName Bool
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> Map VName Bool
mustBeExplicitAux TypeBase (DimDecl VName) ()
t

-- | Figure out which of the sizes in a binding type must be passed
-- explicitly, because their first use is as something else than just
-- an array dimension.
mustBeExplicit :: StructType -> S.Set VName
mustBeExplicit :: TypeBase (DimDecl VName) () -> Set VName
mustBeExplicit TypeBase (DimDecl VName) ()
bind_t =
  let ([TypeBase (DimDecl VName) ()]
ts, TypeBase (DimDecl VName) ()
ret) = TypeBase (DimDecl VName) ()
-> ([TypeBase (DimDecl VName) ()], TypeBase (DimDecl VName) ())
forall dim as.
TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType TypeBase (DimDecl VName) ()
bind_t
      alsoRet :: Map VName Bool -> Map VName Bool
alsoRet =
        (Bool -> Bool -> Bool)
-> Map VName Bool -> Map VName Bool -> Map VName Bool
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Bool -> Bool -> Bool
(&&) (Map VName Bool -> Map VName Bool -> Map VName Bool)
-> Map VName Bool -> Map VName Bool -> Map VName Bool
forall a b. (a -> b) -> a -> b
$
          [(VName, Bool)] -> Map VName Bool
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Bool)] -> Map VName Bool)
-> [(VName, Bool)] -> Map VName Bool
forall a b. (a -> b) -> a -> b
$ [VName] -> [Bool] -> [(VName, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [VName]
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames TypeBase (DimDecl VName) ()
ret) ([Bool] -> [(VName, Bool)]) -> [Bool] -> [(VName, Bool)]
forall a b. (a -> b) -> a -> b
$ Bool -> [Bool]
forall a. a -> [a]
repeat Bool
True
   in [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ Map VName Bool -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName Bool -> [VName]) -> Map VName Bool -> [VName]
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> Map VName Bool -> Map VName Bool
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter Bool -> Bool
forall a. a -> a
id (Map VName Bool -> Map VName Bool)
-> Map VName Bool -> Map VName Bool
forall a b. (a -> b) -> a -> b
$ Map VName Bool -> Map VName Bool
alsoRet (Map VName Bool -> Map VName Bool)
-> Map VName Bool -> Map VName Bool
forall a b. (a -> b) -> a -> b
$ (Map VName Bool -> TypeBase (DimDecl VName) () -> Map VName Bool)
-> Map VName Bool
-> [TypeBase (DimDecl VName) ()]
-> Map VName Bool
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map VName Bool -> TypeBase (DimDecl VName) () -> Map VName Bool
onType Map VName Bool
forall a. Monoid a => a
mempty [TypeBase (DimDecl VName) ()]
ts
  where
    onType :: Map VName Bool -> TypeBase (DimDecl VName) () -> Map VName Bool
onType Map VName Bool
uses TypeBase (DimDecl VName) ()
t = Map VName Bool
uses Map VName Bool -> Map VName Bool -> Map VName Bool
forall a. Semigroup a => a -> a -> a
<> TypeBase (DimDecl VName) () -> Map VName Bool
mustBeExplicitAux TypeBase (DimDecl VName) ()
t -- Left-biased union.

-- | Return the uniqueness of a type.
uniqueness :: TypeBase shape as -> Uniqueness
uniqueness :: TypeBase shape as -> Uniqueness
uniqueness (Array as
_ Uniqueness
u ScalarTypeBase shape ()
_ ShapeDecl shape
_) = Uniqueness
u
uniqueness (Scalar (TypeVar as
_ Uniqueness
u TypeName
_ [TypeArg shape]
_)) = Uniqueness
u
uniqueness (Scalar (Sum Map Name [TypeBase shape as]
ts)) = ([TypeBase shape as] -> Uniqueness)
-> [[TypeBase shape as]] -> Uniqueness
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((TypeBase shape as -> Uniqueness)
-> [TypeBase shape as] -> Uniqueness
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase shape as -> Uniqueness
forall shape as. TypeBase shape as -> Uniqueness
uniqueness) ([[TypeBase shape as]] -> Uniqueness)
-> [[TypeBase shape as]] -> Uniqueness
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase shape as] -> [[TypeBase shape as]]
forall k a. Map k a -> [a]
M.elems Map Name [TypeBase shape as]
ts
uniqueness (Scalar (Record Map Name (TypeBase shape as)
fs)) = (TypeBase shape as -> Uniqueness)
-> [TypeBase shape as] -> Uniqueness
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase shape as -> Uniqueness
forall shape as. TypeBase shape as -> Uniqueness
uniqueness ([TypeBase shape as] -> Uniqueness)
-> [TypeBase shape as] -> Uniqueness
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase shape as) -> [TypeBase shape as]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase shape as)
fs
uniqueness TypeBase shape as
_ = Uniqueness
Nonunique

-- | @unique t@ is 'True' if the type of the argument is unique.
unique :: TypeBase shape as -> Bool
unique :: TypeBase shape as -> Bool
unique = (Uniqueness -> Uniqueness -> Bool
forall a. Eq a => a -> a -> Bool
== Uniqueness
Unique) (Uniqueness -> Bool)
-> (TypeBase shape as -> Uniqueness) -> TypeBase shape as -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase shape as -> Uniqueness
forall shape as. TypeBase shape as -> Uniqueness
uniqueness

-- | Return the set of all variables mentioned in the aliasing of a
-- type.
aliases :: Monoid as => TypeBase shape as -> as
aliases :: TypeBase shape as -> as
aliases = (shape -> as) -> (as -> as) -> TypeBase shape as -> as
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap (as -> shape -> as
forall a b. a -> b -> a
const as
forall a. Monoid a => a
mempty) as -> as
forall a. a -> a
id

-- | @diet t@ returns a description of how a function parameter of
-- type @t@ might consume its argument.
diet :: TypeBase shape as -> Diet
diet :: TypeBase shape as -> Diet
diet (Scalar (Record Map Name (TypeBase shape as)
ets)) = Map Name Diet -> Diet
RecordDiet (Map Name Diet -> Diet) -> Map Name Diet -> Diet
forall a b. (a -> b) -> a -> b
$ (TypeBase shape as -> Diet)
-> Map Name (TypeBase shape as) -> Map Name Diet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape as -> Diet
forall shape as. TypeBase shape as -> Diet
diet Map Name (TypeBase shape as)
ets
diet (Scalar (Prim PrimType
_)) = Diet
Observe
diet (Scalar (Arrow as
_ PName
_ TypeBase shape as
t1 TypeBase shape as
t2)) = Diet -> Diet -> Diet
FuncDiet (TypeBase shape as -> Diet
forall shape as. TypeBase shape as -> Diet
diet TypeBase shape as
t1) (TypeBase shape as -> Diet
forall shape as. TypeBase shape as -> Diet
diet TypeBase shape as
t2)
diet (Array as
_ Uniqueness
Unique ScalarTypeBase shape ()
_ ShapeDecl shape
_) = Diet
Consume
diet (Array as
_ Uniqueness
Nonunique ScalarTypeBase shape ()
_ ShapeDecl shape
_) = Diet
Observe
diet (Scalar (TypeVar as
_ Uniqueness
Unique TypeName
_ [TypeArg shape]
_)) = Diet
Consume
diet (Scalar (TypeVar as
_ Uniqueness
Nonunique TypeName
_ [TypeArg shape]
_)) = Diet
Observe
diet (Scalar Sum {}) = Diet
Observe

-- | Convert any type to one that has rank information, no alias
-- information, and no embedded names.
toStructural ::
  TypeBase dim as ->
  TypeBase () ()
toStructural :: TypeBase dim as -> TypeBase () ()
toStructural = (TypeBase () as -> () -> TypeBase () ())
-> () -> TypeBase () as -> TypeBase () ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip TypeBase () as -> () -> TypeBase () ()
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases () (TypeBase () as -> TypeBase () ())
-> (TypeBase dim as -> TypeBase () as)
-> TypeBase dim as
-> TypeBase () ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (dim -> ()) -> TypeBase dim as -> TypeBase () as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> dim -> ()
forall a b. a -> b -> a
const ())

-- | Remove aliasing information from a type.
toStruct ::
  TypeBase dim as ->
  TypeBase dim ()
toStruct :: TypeBase dim as -> TypeBase dim ()
toStruct TypeBase dim as
t = TypeBase dim as
t TypeBase dim as -> () -> TypeBase dim ()
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()

-- | Replace no aliasing with an empty alias set.
fromStruct ::
  TypeBase dim as ->
  TypeBase dim Aliasing
fromStruct :: TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase dim as
t = TypeBase dim as
t TypeBase dim as -> Aliasing -> TypeBase dim Aliasing
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliasing
forall a. Set a
S.empty

-- | @peelArray n t@ returns the type resulting from peeling the first
-- @n@ array dimensions from @t@.  Returns @Nothing@ if @t@ has less
-- than @n@ dimensions.
peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray Int
n (Array as
als Uniqueness
u ScalarTypeBase dim ()
t ShapeDecl dim
shape)
  | ShapeDecl dim -> Int
forall dim. ShapeDecl dim -> Int
shapeRank ShapeDecl dim
shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n =
    TypeBase dim as -> Maybe (TypeBase dim as)
forall a. a -> Maybe a
Just (TypeBase dim as -> Maybe (TypeBase dim as))
-> TypeBase dim as -> Maybe (TypeBase dim as)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
t TypeBase dim () -> (() -> as) -> TypeBase dim as
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` as -> () -> as
forall a b. a -> b -> a
const as
als
  | Bool
otherwise =
    as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array as
als Uniqueness
u ScalarTypeBase dim ()
t (ShapeDecl dim -> TypeBase dim as)
-> Maybe (ShapeDecl dim) -> Maybe (TypeBase dim as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> ShapeDecl dim -> Maybe (ShapeDecl dim)
forall dim. Int -> ShapeDecl dim -> Maybe (ShapeDecl dim)
stripDims Int
n ShapeDecl dim
shape
peelArray Int
_ TypeBase dim as
_ = Maybe (TypeBase dim as)
forall a. Maybe a
Nothing

-- | @arrayOf t s u@ constructs an array type.  The convenience
-- compared to using the 'Array' constructor directly is that @t@ can
-- itself be an array.  If @t@ is an @n@-dimensional array, and @s@ is
-- a list of length @n@, the resulting type is of an @n+m@ dimensions.
-- The uniqueness of the new array will be @u@, no matter the
-- uniqueness of @t@.
arrayOf ::
  Monoid as =>
  TypeBase dim as ->
  ShapeDecl dim ->
  Uniqueness ->
  TypeBase dim as
arrayOf :: TypeBase dim as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOf TypeBase dim as
t = TypeBase dim as
-> as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
forall as dim.
Monoid as =>
TypeBase dim as
-> as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOfWithAliases (TypeBase dim as
t TypeBase dim as -> Uniqueness -> TypeBase dim as
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique) as
forall a. Monoid a => a
mempty

arrayOfWithAliases ::
  Monoid as =>
  TypeBase dim as ->
  as ->
  ShapeDecl dim ->
  Uniqueness ->
  TypeBase dim as
arrayOfWithAliases :: TypeBase dim as
-> as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOfWithAliases (Array as
as1 Uniqueness
_ ScalarTypeBase dim ()
et ShapeDecl dim
shape1) as
as2 ShapeDecl dim
shape2 Uniqueness
u =
  as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array (as
as1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
as2) Uniqueness
u ScalarTypeBase dim ()
et (ShapeDecl dim
shape2 ShapeDecl dim -> ShapeDecl dim -> ShapeDecl dim
forall a. Semigroup a => a -> a -> a
<> ShapeDecl dim
shape1)
arrayOfWithAliases (Scalar ScalarTypeBase dim as
t) as
as ShapeDecl dim
shape Uniqueness
u =
  as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array as
as Uniqueness
u ((as -> ()) -> ScalarTypeBase dim as -> ScalarTypeBase dim ()
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (() -> as -> ()
forall a b. a -> b -> a
const ()) ScalarTypeBase dim as
t) ShapeDecl dim
shape

-- | @stripArray n t@ removes the @n@ outermost layers of the array.
-- Essentially, it is the type of indexing an array of type @t@ with
-- @n@ indexes.
stripArray :: Int -> TypeBase dim as -> TypeBase dim as
stripArray :: Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
n (Array as
als Uniqueness
u ScalarTypeBase dim ()
et ShapeDecl dim
shape)
  | Just ShapeDecl dim
shape' <- Int -> ShapeDecl dim -> Maybe (ShapeDecl dim)
forall dim. Int -> ShapeDecl dim -> Maybe (ShapeDecl dim)
stripDims Int
n ShapeDecl dim
shape =
    as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array as
als Uniqueness
u ScalarTypeBase dim ()
et ShapeDecl dim
shape'
  | Bool
otherwise =
    ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
et TypeBase dim () -> Uniqueness -> TypeBase dim ()
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u TypeBase dim () -> as -> TypeBase dim as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` as
als
stripArray Int
_ TypeBase dim as
t = TypeBase dim as
t

-- | Create a record type corresponding to a tuple with the given
-- element types.
tupleRecord :: [TypeBase dim as] -> TypeBase dim as
tupleRecord :: [TypeBase dim as] -> TypeBase dim as
tupleRecord = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ([TypeBase dim as] -> ScalarTypeBase dim as)
-> [TypeBase dim as]
-> TypeBase dim as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim as) -> ScalarTypeBase dim as)
-> ([TypeBase dim as] -> Map Name (TypeBase dim as))
-> [TypeBase dim as]
-> ScalarTypeBase dim as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, TypeBase dim as)] -> Map Name (TypeBase dim as)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase dim as)] -> Map Name (TypeBase dim as))
-> ([TypeBase dim as] -> [(Name, TypeBase dim as)])
-> [TypeBase dim as]
-> Map Name (TypeBase dim as)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Name] -> [TypeBase dim as] -> [(Name, TypeBase dim as)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames

-- | Does this type corespond to a tuple?  If so, return the elements
-- of that tuple.
isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord (Scalar (Record Map Name (TypeBase dim as)
fs)) = Map Name (TypeBase dim as) -> Maybe [TypeBase dim as]
forall a. Map Name a -> Maybe [a]
areTupleFields Map Name (TypeBase dim as)
fs
isTupleRecord TypeBase dim as
_ = Maybe [TypeBase dim as]
forall a. Maybe a
Nothing

-- | Does this record map correspond to a tuple?
areTupleFields :: M.Map Name a -> Maybe [a]
areTupleFields :: Map Name a -> Maybe [a]
areTupleFields Map Name a
fs =
  let fs' :: [(Name, a)]
fs' = Map Name a -> [(Name, a)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name a
fs
   in if [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Name -> Name -> Bool) -> [Name] -> [Name] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
(==) (((Name, a) -> Name) -> [(Name, a)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, a) -> Name
forall a b. (a, b) -> a
fst [(Name, a)]
fs') [Name]
tupleFieldNames
        then [a] -> Maybe [a]
forall a. a -> Maybe a
Just ([a] -> Maybe [a]) -> [a] -> Maybe [a]
forall a b. (a -> b) -> a -> b
$ ((Name, a) -> a) -> [(Name, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Name, a) -> a
forall a b. (a, b) -> b
snd [(Name, a)]
fs'
        else Maybe [a]
forall a. Maybe a
Nothing

-- | Construct a record map corresponding to a tuple.
tupleFields :: [a] -> M.Map Name a
tupleFields :: [a] -> Map Name a
tupleFields [a]
as = [(Name, a)] -> Map Name a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, a)] -> Map Name a) -> [(Name, a)] -> Map Name a
forall a b. (a -> b) -> a -> b
$ [Name] -> [a] -> [(Name, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [a]
as

-- | Increasing field names for a tuple (starts at 0).
tupleFieldNames :: [Name]
tupleFieldNames :: [Name]
tupleFieldNames = (Int -> Name) -> [Int] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Name
nameFromString (String -> Name) -> (Int -> String) -> Int -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show) [(Int
0 :: Int) ..]

-- | Sort fields by their name; taking care to sort numeric fields by
-- their numeric value.  This ensures that tuples and tuple-like
-- records match.
sortFields :: M.Map Name a -> [(Name, a)]
sortFields :: Map Name a -> [(Name, a)]
sortFields Map Name a
l = ((Either Int Name, (Name, a)) -> (Name, a))
-> [(Either Int Name, (Name, a))] -> [(Name, a)]
forall a b. (a -> b) -> [a] -> [b]
map (Either Int Name, (Name, a)) -> (Name, a)
forall a b. (a, b) -> b
snd ([(Either Int Name, (Name, a))] -> [(Name, a)])
-> [(Either Int Name, (Name, a))] -> [(Name, a)]
forall a b. (a -> b) -> a -> b
$ ((Either Int Name, (Name, a)) -> Either Int Name)
-> [(Either Int Name, (Name, a))] -> [(Either Int Name, (Name, a))]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Either Int Name, (Name, a)) -> Either Int Name
forall a b. (a, b) -> a
fst ([(Either Int Name, (Name, a))] -> [(Either Int Name, (Name, a))])
-> [(Either Int Name, (Name, a))] -> [(Either Int Name, (Name, a))]
forall a b. (a -> b) -> a -> b
$ [Either Int Name] -> [(Name, a)] -> [(Either Int Name, (Name, a))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, a) -> Either Int Name) -> [(Name, a)] -> [Either Int Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Either Int Name
fieldish (Name -> Either Int Name)
-> ((Name, a) -> Name) -> (Name, a) -> Either Int Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, a) -> Name
forall a b. (a, b) -> a
fst) [(Name, a)]
l') [(Name, a)]
l'
  where
    l' :: [(Name, a)]
l' = Map Name a -> [(Name, a)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name a
l
    fieldish :: Name -> Either Int Name
fieldish Name
s = case ReadS Int
forall a. Read a => ReadS a
reads ReadS Int -> ReadS Int
forall a b. (a -> b) -> a -> b
$ Name -> String
nameToString Name
s of
      [(Int
x, String
"")] -> Int -> Either Int Name
forall a b. a -> Either a b
Left (Int
x :: Int)
      [(Int, String)]
_ -> Name -> Either Int Name
forall a b. b -> Either a b
Right Name
s

-- | Sort the constructors of a sum type in some well-defined (but not
-- otherwise significant) manner.
sortConstrs :: M.Map Name a -> [(Name, a)]
sortConstrs :: Map Name a -> [(Name, a)]
sortConstrs Map Name a
cs = ((Name, a) -> Name) -> [(Name, a)] -> [(Name, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, a) -> Name
forall a b. (a, b) -> a
fst ([(Name, a)] -> [(Name, a)]) -> [(Name, a)] -> [(Name, a)]
forall a b. (a -> b) -> a -> b
$ Map Name a -> [(Name, a)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name a
cs

-- | Is this a 'TypeParamType'?
isTypeParam :: TypeParamBase vn -> Bool
isTypeParam :: TypeParamBase vn -> Bool
isTypeParam TypeParamType {} = Bool
True
isTypeParam TypeParamDim {} = Bool
False

-- | Is this a 'TypeParamDim'?
isSizeParam :: TypeParamBase vn -> Bool
isSizeParam :: TypeParamBase vn -> Bool
isSizeParam = Bool -> Bool
not (Bool -> Bool)
-> (TypeParamBase vn -> Bool) -> TypeParamBase vn -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase vn -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam

-- | Combine the shape information of types as much as possible. The first
-- argument is the orignal type and the second is the type of the transformed
-- expression. This is necessary since the original type may contain additional
-- information (e.g., shape restrictions) from the user given annotation.
combineTypeShapes ::
  (Monoid as, ArrayDim dim) =>
  TypeBase dim as ->
  TypeBase dim as ->
  TypeBase dim as
combineTypeShapes :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes (Scalar (Record Map Name (TypeBase dim as)
ts1)) (Scalar (Record Map Name (TypeBase dim as)
ts2))
  | Map Name (TypeBase dim as) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase dim as)
ts1 [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name (TypeBase dim as) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase dim as)
ts2 =
    ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$
      Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim as) -> ScalarTypeBase dim as)
-> Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$
        ((TypeBase dim as, TypeBase dim as) -> TypeBase dim as)
-> Map Name (TypeBase dim as, TypeBase dim as)
-> Map Name (TypeBase dim as)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map
          ((TypeBase dim as -> TypeBase dim as -> TypeBase dim as)
-> (TypeBase dim as, TypeBase dim as) -> TypeBase dim as
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TypeBase dim as -> TypeBase dim as -> TypeBase dim as
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes)
          ((TypeBase dim as
 -> TypeBase dim as -> (TypeBase dim as, TypeBase dim as))
-> Map Name (TypeBase dim as)
-> Map Name (TypeBase dim as)
-> Map Name (TypeBase dim as, TypeBase dim as)
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 as)
ts1 Map Name (TypeBase dim as)
ts2)
combineTypeShapes (Scalar (Sum Map Name [TypeBase dim as]
cs1)) (Scalar (Sum Map Name [TypeBase dim as]
cs2))
  | Map Name [TypeBase dim as] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase dim as]
cs1 [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name [TypeBase dim as] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase dim as]
cs2 =
    ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$
      Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim as] -> ScalarTypeBase dim as)
-> Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$
        (([TypeBase dim as], [TypeBase dim as]) -> [TypeBase dim as])
-> Map Name ([TypeBase dim as], [TypeBase dim as])
-> Map Name [TypeBase dim as]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map
          (([TypeBase dim as] -> [TypeBase dim as] -> [TypeBase dim as])
-> ([TypeBase dim as], [TypeBase dim as]) -> [TypeBase dim as]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (([TypeBase dim as] -> [TypeBase dim as] -> [TypeBase dim as])
 -> ([TypeBase dim as], [TypeBase dim as]) -> [TypeBase dim as])
-> ([TypeBase dim as] -> [TypeBase dim as] -> [TypeBase dim as])
-> ([TypeBase dim as], [TypeBase dim as])
-> [TypeBase dim as]
forall a b. (a -> b) -> a -> b
$ (TypeBase dim as -> TypeBase dim as -> TypeBase dim as)
-> [TypeBase dim as] -> [TypeBase dim as] -> [TypeBase dim as]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase dim as -> TypeBase dim as -> TypeBase dim as
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes)
          (([TypeBase dim as]
 -> [TypeBase dim as] -> ([TypeBase dim as], [TypeBase dim as]))
-> Map Name [TypeBase dim as]
-> Map Name [TypeBase dim as]
-> Map Name ([TypeBase dim as], [TypeBase dim as])
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 as]
cs1 Map Name [TypeBase dim as]
cs2)
combineTypeShapes (Scalar (Arrow as
als1 PName
p1 TypeBase dim as
a1 TypeBase dim as
b1)) (Scalar (Arrow as
als2 PName
_p2 TypeBase dim as
a2 TypeBase dim as
b2)) =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow (as
als1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als2) PName
p1 (TypeBase dim as -> TypeBase dim as -> TypeBase dim as
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes TypeBase dim as
a1 TypeBase dim as
a2) (TypeBase dim as -> TypeBase dim as -> TypeBase dim as
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes TypeBase dim as
b1 TypeBase dim as
b2)
combineTypeShapes (Scalar (TypeVar as
als1 Uniqueness
u1 TypeName
v [TypeArg dim]
targs1)) (Scalar (TypeVar as
als2 Uniqueness
_ TypeName
_ [TypeArg dim]
targs2)) =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar (as
als1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als2) Uniqueness
u1 TypeName
v ([TypeArg dim] -> ScalarTypeBase dim as)
-> [TypeArg dim] -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ (TypeArg dim -> TypeArg dim -> TypeArg dim)
-> [TypeArg dim] -> [TypeArg dim] -> [TypeArg dim]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeArg dim -> TypeArg dim -> TypeArg dim
forall dim.
ArrayDim dim =>
TypeArg dim -> TypeArg dim -> TypeArg dim
f [TypeArg dim]
targs1 [TypeArg dim]
targs2
  where
    f :: TypeArg dim -> TypeArg dim -> TypeArg dim
f (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 () -> TypeBase dim () -> TypeBase dim ()
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes TypeBase dim ()
t1 TypeBase dim ()
t2) SrcLoc
loc
    f TypeArg dim
targ TypeArg dim
_ = TypeArg dim
targ
combineTypeShapes (Array as
als1 Uniqueness
u1 ScalarTypeBase dim ()
et1 ShapeDecl dim
shape1) (Array as
als2 Uniqueness
_u2 ScalarTypeBase dim ()
et2 ShapeDecl dim
_shape2) =
  TypeBase dim as
-> as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
forall as dim.
Monoid as =>
TypeBase dim as
-> as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOfWithAliases
    ( TypeBase dim () -> TypeBase dim () -> TypeBase dim ()
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes (ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
et1) (ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
et2)
        TypeBase dim () -> as -> TypeBase dim as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` as
forall a. Monoid a => a
mempty
    )
    (as
als1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als2)
    ShapeDecl dim
shape1
    Uniqueness
u1
combineTypeShapes TypeBase dim as
_ TypeBase dim as
new_tp = TypeBase dim as
new_tp

-- | Match the dimensions of otherwise assumed-equal types.
matchDims ::
  (Monoid as, Monad m) =>
  (d1 -> d2 -> m d1) ->
  TypeBase d1 as ->
  TypeBase d2 as ->
  m (TypeBase d1 as)
matchDims :: (d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims d1 -> d2 -> m d1
onDims TypeBase d1 as
t1 TypeBase d2 as
t2 =
  case (TypeBase d1 as
t1, TypeBase d2 as
t2) of
    (Array as
als1 Uniqueness
u1 ScalarTypeBase d1 ()
et1 ShapeDecl d1
shape1, Array as
als2 Uniqueness
u2 ScalarTypeBase d2 ()
et2 ShapeDecl d2
shape2) ->
      (TypeBase d1 () -> as -> TypeBase d1 as)
-> as -> TypeBase d1 () -> TypeBase d1 as
forall a b c. (a -> b -> c) -> b -> a -> c
flip TypeBase d1 () -> as -> TypeBase d1 as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases (as
als1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als2)
        (TypeBase d1 () -> TypeBase d1 as)
-> m (TypeBase d1 ()) -> m (TypeBase d1 as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( TypeBase d1 () -> ShapeDecl d1 -> Uniqueness -> TypeBase d1 ()
forall as dim.
Monoid as =>
TypeBase dim as -> ShapeDecl dim -> Uniqueness -> TypeBase dim as
arrayOf
                (TypeBase d1 () -> ShapeDecl d1 -> Uniqueness -> TypeBase d1 ())
-> m (TypeBase d1 ())
-> m (ShapeDecl d1 -> Uniqueness -> TypeBase d1 ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (d1 -> d2 -> m d1)
-> TypeBase d1 () -> TypeBase d2 () -> m (TypeBase d1 ())
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims d1 -> d2 -> m d1
onDims (ScalarTypeBase d1 () -> TypeBase d1 ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d1 ()
et1) (ScalarTypeBase d2 () -> TypeBase d2 ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d2 ()
et2)
                m (ShapeDecl d1 -> Uniqueness -> TypeBase d1 ())
-> m (ShapeDecl d1) -> m (Uniqueness -> TypeBase d1 ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ShapeDecl d1 -> ShapeDecl d2 -> m (ShapeDecl d1)
onShapes ShapeDecl d1
shape1 ShapeDecl d2
shape2
                m (Uniqueness -> TypeBase d1 ())
-> m Uniqueness -> m (TypeBase d1 ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Uniqueness -> m Uniqueness
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Uniqueness -> Uniqueness -> Uniqueness
forall a. Ord a => a -> a -> a
min Uniqueness
u1 Uniqueness
u2)
            )
    (Scalar (Record Map Name (TypeBase d1 as)
f1), Scalar (Record Map Name (TypeBase d2 as)
f2)) ->
      ScalarTypeBase d1 as -> TypeBase d1 as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase d1 as -> TypeBase d1 as)
-> (Map Name (TypeBase d1 as) -> ScalarTypeBase d1 as)
-> Map Name (TypeBase d1 as)
-> TypeBase d1 as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase d1 as) -> ScalarTypeBase d1 as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record
        (Map Name (TypeBase d1 as) -> TypeBase d1 as)
-> m (Map Name (TypeBase d1 as)) -> m (TypeBase d1 as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((TypeBase d1 as, TypeBase d2 as) -> m (TypeBase d1 as))
-> Map Name (TypeBase d1 as, TypeBase d2 as)
-> m (Map Name (TypeBase d1 as))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as))
-> (TypeBase d1 as, TypeBase d2 as) -> m (TypeBase d1 as)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims d1 -> d2 -> m d1
onDims)) ((TypeBase d1 as
 -> TypeBase d2 as -> (TypeBase d1 as, TypeBase d2 as))
-> Map Name (TypeBase d1 as)
-> Map Name (TypeBase d2 as)
-> Map Name (TypeBase d1 as, TypeBase d2 as)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name (TypeBase d1 as)
f1 Map Name (TypeBase d2 as)
f2)
    (Scalar (Sum Map Name [TypeBase d1 as]
cs1), Scalar (Sum Map Name [TypeBase d2 as]
cs2)) ->
      ScalarTypeBase d1 as -> TypeBase d1 as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase d1 as -> TypeBase d1 as)
-> (Map Name [TypeBase d1 as] -> ScalarTypeBase d1 as)
-> Map Name [TypeBase d1 as]
-> TypeBase d1 as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase d1 as] -> ScalarTypeBase d1 as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum
        (Map Name [TypeBase d1 as] -> TypeBase d1 as)
-> m (Map Name [TypeBase d1 as]) -> m (TypeBase d1 as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(TypeBase d1 as, TypeBase d2 as)] -> m [TypeBase d1 as])
-> Map Name [(TypeBase d1 as, TypeBase d2 as)]
-> m (Map Name [TypeBase d1 as])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
          (((TypeBase d1 as, TypeBase d2 as) -> m (TypeBase d1 as))
-> [(TypeBase d1 as, TypeBase d2 as)] -> m [TypeBase d1 as]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as))
-> (TypeBase d1 as, TypeBase d2 as) -> m (TypeBase d1 as)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims d1 -> d2 -> m d1
onDims)))
          (([TypeBase d1 as]
 -> [TypeBase d2 as] -> [(TypeBase d1 as, TypeBase d2 as)])
-> Map Name [TypeBase d1 as]
-> Map Name [TypeBase d2 as]
-> Map Name [(TypeBase d1 as, TypeBase d2 as)]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith [TypeBase d1 as]
-> [TypeBase d2 as] -> [(TypeBase d1 as, TypeBase d2 as)]
forall a b. [a] -> [b] -> [(a, b)]
zip Map Name [TypeBase d1 as]
cs1 Map Name [TypeBase d2 as]
cs2)
    (Scalar (Arrow as
als1 PName
p1 TypeBase d1 as
a1 TypeBase d1 as
b1), Scalar (Arrow as
als2 PName
_p2 TypeBase d2 as
a2 TypeBase d2 as
b2)) ->
      ScalarTypeBase d1 as -> TypeBase d1 as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar
        (ScalarTypeBase d1 as -> TypeBase d1 as)
-> m (ScalarTypeBase d1 as) -> m (TypeBase d1 as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (as
-> PName
-> TypeBase d1 as
-> TypeBase d1 as
-> ScalarTypeBase d1 as
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow (as
als1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als2) PName
p1 (TypeBase d1 as -> TypeBase d1 as -> ScalarTypeBase d1 as)
-> m (TypeBase d1 as) -> m (TypeBase d1 as -> ScalarTypeBase d1 as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims d1 -> d2 -> m d1
onDims TypeBase d1 as
a1 TypeBase d2 as
a2 m (TypeBase d1 as -> ScalarTypeBase d1 as)
-> m (TypeBase d1 as) -> m (ScalarTypeBase d1 as)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims d1 -> d2 -> m d1
onDims TypeBase d1 as
b1 TypeBase d2 as
b2)
    ( Scalar (TypeVar as
als1 Uniqueness
u TypeName
v [TypeArg d1]
targs1),
      Scalar (TypeVar as
als2 Uniqueness
_ TypeName
_ [TypeArg d2]
targs2)
      ) ->
        ScalarTypeBase d1 as -> TypeBase d1 as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase d1 as -> TypeBase d1 as)
-> ([TypeArg d1] -> ScalarTypeBase d1 as)
-> [TypeArg d1]
-> TypeBase d1 as
forall b c a. (b -> c) -> (a -> b) -> a -> c
. as
-> Uniqueness -> TypeName -> [TypeArg d1] -> ScalarTypeBase d1 as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar (as
als1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
als2) Uniqueness
u TypeName
v ([TypeArg d1] -> TypeBase d1 as)
-> m [TypeArg d1] -> m (TypeBase d1 as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg d1 -> TypeArg d2 -> m (TypeArg d1))
-> [TypeArg d1] -> [TypeArg d2] -> m [TypeArg d1]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeArg d1 -> TypeArg d2 -> m (TypeArg d1)
forall (m :: * -> *) dim p.
Monad m =>
TypeArg dim -> p -> m (TypeArg dim)
matchTypeArg [TypeArg d1]
targs1 [TypeArg d2]
targs2
    (TypeBase d1 as, TypeBase d2 as)
_ -> TypeBase d1 as -> m (TypeBase d1 as)
forall (m :: * -> *) a. Monad m => a -> m a
return TypeBase d1 as
t1
  where
    matchTypeArg :: TypeArg dim -> p -> m (TypeArg dim)
matchTypeArg ta :: TypeArg dim
ta@TypeArgType {} p
_ = TypeArg dim -> m (TypeArg dim)
forall (m :: * -> *) a. Monad m => a -> m a
return TypeArg dim
ta
    matchTypeArg TypeArg dim
a p
_ = TypeArg dim -> m (TypeArg dim)
forall (m :: * -> *) a. Monad m => a -> m a
return TypeArg dim
a

    onShapes :: ShapeDecl d1 -> ShapeDecl d2 -> m (ShapeDecl d1)
onShapes ShapeDecl d1
shape1 ShapeDecl d2
shape2 =
      [d1] -> ShapeDecl d1
forall dim. [dim] -> ShapeDecl dim
ShapeDecl ([d1] -> ShapeDecl d1) -> m [d1] -> m (ShapeDecl d1)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (d1 -> d2 -> m d1) -> [d1] -> [d2] -> m [d1]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM d1 -> d2 -> m d1
onDims (ShapeDecl d1 -> [d1]
forall dim. ShapeDecl dim -> [dim]
shapeDims ShapeDecl d1
shape1) (ShapeDecl d2 -> [d2]
forall dim. ShapeDecl dim -> [dim]
shapeDims ShapeDecl d2
shape2)

-- | Set the uniqueness attribute of a type.  If the type is a record
-- or sum type, the uniqueness of its components will be modified.
setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as
setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as
setUniqueness (Array as
als Uniqueness
_ ScalarTypeBase dim ()
et ShapeDecl dim
shape) Uniqueness
u =
  as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array as
als Uniqueness
u ScalarTypeBase dim ()
et ShapeDecl dim
shape
setUniqueness (Scalar (TypeVar as
als Uniqueness
_ TypeName
t [TypeArg dim]
targs)) Uniqueness
u =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar as
als Uniqueness
u TypeName
t [TypeArg dim]
targs
setUniqueness (Scalar (Record Map Name (TypeBase dim as)
ets)) Uniqueness
u =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim as) -> ScalarTypeBase dim as)
-> Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ (TypeBase dim as -> TypeBase dim as)
-> Map Name (TypeBase dim as) -> Map Name (TypeBase dim as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (TypeBase dim as -> Uniqueness -> TypeBase dim as
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u) Map Name (TypeBase dim as)
ets
setUniqueness (Scalar (Sum Map Name [TypeBase dim as]
ets)) Uniqueness
u =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim as] -> ScalarTypeBase dim as)
-> Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim as] -> [TypeBase dim as])
-> Map Name [TypeBase dim as] -> Map Name [TypeBase dim as]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim as -> TypeBase dim as)
-> [TypeBase dim as] -> [TypeBase dim as]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase dim as -> Uniqueness -> TypeBase dim as
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u)) Map Name [TypeBase dim as]
ets
setUniqueness TypeBase dim as
t Uniqueness
_ = TypeBase dim as
t

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

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

intValueType :: IntValue -> IntType
intValueType :: IntValue -> IntType
intValueType Int8Value {} = IntType
Int8
intValueType Int16Value {} = IntType
Int16
intValueType Int32Value {} = IntType
Int32
intValueType Int64Value {} = IntType
Int64

floatValueType :: FloatValue -> FloatType
floatValueType :: FloatValue -> FloatType
floatValueType Float32Value {} = FloatType
Float32
floatValueType Float64Value {} = FloatType
Float64

-- | The type of a basic value.
primValueType :: PrimValue -> PrimType
primValueType :: PrimValue -> PrimType
primValueType (SignedValue IntValue
v) = IntType -> PrimType
Signed (IntType -> PrimType) -> IntType -> PrimType
forall a b. (a -> b) -> a -> b
$ IntValue -> IntType
intValueType IntValue
v
primValueType (UnsignedValue IntValue
v) = IntType -> PrimType
Unsigned (IntType -> PrimType) -> IntType -> PrimType
forall a b. (a -> b) -> a -> b
$ IntValue -> IntType
intValueType IntValue
v
primValueType (FloatValue FloatValue
v) = FloatType -> PrimType
FloatType (FloatType -> PrimType) -> FloatType -> PrimType
forall a b. (a -> b) -> a -> b
$ FloatValue -> FloatType
floatValueType FloatValue
v
primValueType BoolValue {} = PrimType
Bool

-- | The type of the value.
valueType :: Value -> ValueType
valueType :: Value -> ValueType
valueType (PrimValue PrimValue
bv) = ScalarTypeBase Int64 () -> ValueType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Int64 () -> ValueType)
-> ScalarTypeBase Int64 () -> ValueType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Int64 ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase Int64 ())
-> PrimType -> ScalarTypeBase Int64 ()
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
bv
valueType (ArrayValue Array Int Value
_ ValueType
t) = ValueType
t

-- | The size of values of this type, in bytes.
primByteSize :: Num a => PrimType -> a
primByteSize :: PrimType -> a
primByteSize (Signed IntType
it) = IntType -> a
forall a. Num a => IntType -> a
Primitive.intByteSize IntType
it
primByteSize (Unsigned IntType
it) = IntType -> a
forall a. Num a => IntType -> a
Primitive.intByteSize IntType
it
primByteSize (FloatType FloatType
ft) = FloatType -> a
forall a. Num a => FloatType -> a
Primitive.floatByteSize FloatType
ft
primByteSize PrimType
Bool = a
1

-- | Construct a 'ShapeDecl' with the given number of 'AnyDim'
-- dimensions.
rank :: Int -> ShapeDecl (DimDecl VName)
rank :: Int -> ShapeDecl (DimDecl VName)
rank Int
n = [DimDecl VName] -> ShapeDecl (DimDecl VName)
forall dim. [dim] -> ShapeDecl dim
ShapeDecl ([DimDecl VName] -> ShapeDecl (DimDecl VName))
-> [DimDecl VName] -> ShapeDecl (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ Int -> DimDecl VName -> [DimDecl VName]
forall a. Int -> a -> [a]
replicate Int
n DimDecl VName
forall vn. DimDecl vn
AnyDim

-- | The type is leaving a scope, so clean up any aliases that
-- reference the bound variables, and turn any dimensions that name
-- them into AnyDim instead.
unscopeType :: S.Set VName -> PatternType -> PatternType
unscopeType :: Set VName -> PatternType -> PatternType
unscopeType Set VName
bound_here PatternType
t = (DimDecl VName -> DimDecl VName) -> PatternType -> PatternType
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
onDim (PatternType -> PatternType) -> PatternType -> PatternType
forall a b. (a -> b) -> a -> b
$ PatternType
t PatternType -> (Aliasing -> Aliasing) -> PatternType
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` (Alias -> Alias) -> Aliasing -> Aliasing
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map Alias -> Alias
unbind
  where
    unbind :: Alias -> Alias
unbind (AliasBound VName
v) | VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound_here = VName -> Alias
AliasFree VName
v
    unbind Alias
a = Alias
a
    onDim :: DimDecl VName -> DimDecl VName
onDim (NamedDim QualName VName
qn) | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound_here = DimDecl VName
forall vn. DimDecl vn
AnyDim
    onDim DimDecl VName
d = DimDecl VName
d

-- | Perform some operation on a given record field.  Returns
-- 'Nothing' if that field does not exist.
onRecordField ::
  (TypeBase dim als -> TypeBase dim als) ->
  [Name] ->
  TypeBase dim als ->
  Maybe (TypeBase dim als)
onRecordField :: (TypeBase dim als -> TypeBase dim als)
-> [Name] -> TypeBase dim als -> Maybe (TypeBase dim als)
onRecordField TypeBase dim als -> TypeBase dim als
f [] TypeBase dim als
t = TypeBase dim als -> Maybe (TypeBase dim als)
forall a. a -> Maybe a
Just (TypeBase dim als -> Maybe (TypeBase dim als))
-> TypeBase dim als -> Maybe (TypeBase dim als)
forall a b. (a -> b) -> a -> b
$ TypeBase dim als -> TypeBase dim als
f TypeBase dim als
t
onRecordField TypeBase dim als -> TypeBase dim als
f (Name
k : [Name]
ks) (Scalar (Record Map Name (TypeBase dim als)
m)) = do
  TypeBase dim als
t <- (TypeBase dim als -> TypeBase dim als)
-> [Name] -> TypeBase dim als -> Maybe (TypeBase dim als)
forall dim als.
(TypeBase dim als -> TypeBase dim als)
-> [Name] -> TypeBase dim als -> Maybe (TypeBase dim als)
onRecordField TypeBase dim als -> TypeBase dim als
f [Name]
ks (TypeBase dim als -> Maybe (TypeBase dim als))
-> Maybe (TypeBase dim als) -> Maybe (TypeBase dim als)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> Map Name (TypeBase dim als) -> Maybe (TypeBase dim als)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
k Map Name (TypeBase dim als)
m
  TypeBase dim als -> Maybe (TypeBase dim als)
forall a. a -> Maybe a
Just (TypeBase dim als -> Maybe (TypeBase dim als))
-> TypeBase dim als -> Maybe (TypeBase dim als)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim als -> TypeBase dim als
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim als -> TypeBase dim als)
-> ScalarTypeBase dim als -> TypeBase dim als
forall a b. (a -> b) -> a -> b
$ 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)
-> Map Name (TypeBase dim als) -> ScalarTypeBase dim als
forall a b. (a -> b) -> a -> b
$ Name
-> TypeBase dim als
-> Map Name (TypeBase dim als)
-> Map Name (TypeBase dim als)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
k TypeBase dim als
t Map Name (TypeBase dim als)
m
onRecordField TypeBase dim als -> TypeBase dim als
_ [Name]
_ TypeBase dim als
_ = Maybe (TypeBase dim als)
forall a. Maybe a
Nothing

-- | The type of an Futhark term.  The aliasing will refer to itself, if
-- the term is a non-tuple-typed variable.
typeOf :: ExpBase Info VName -> PatternType
typeOf :: ExpBase Info VName -> PatternType
typeOf (Literal PrimValue
val SrcLoc
_) = ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) Aliasing)
-> PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
typeOf (IntLit Integer
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (FloatLit Double
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (Parens ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
typeOf (QualParens (QualName VName, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
typeOf (TupLit [ExpBase Info VName]
es SrcLoc
_) = [PatternType] -> PatternType
forall dim as. [TypeBase dim as] -> TypeBase dim as
tupleRecord ([PatternType] -> PatternType) -> [PatternType] -> PatternType
forall a b. (a -> b) -> a -> b
$ (ExpBase Info VName -> PatternType)
-> [ExpBase Info VName] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map ExpBase Info VName -> PatternType
typeOf [ExpBase Info VName]
es
typeOf (RecordLit [FieldBase Info VName]
fs SrcLoc
_) =
  -- Reverse, because M.unions is biased to the left.
  ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ Map Name PatternType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name PatternType -> ScalarTypeBase (DimDecl VName) Aliasing)
-> Map Name PatternType -> ScalarTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ [Map Name PatternType] -> Map Name PatternType
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name PatternType] -> Map Name PatternType)
-> [Map Name PatternType] -> Map Name PatternType
forall a b. (a -> b) -> a -> b
$ [Map Name PatternType] -> [Map Name PatternType]
forall a. [a] -> [a]
reverse ([Map Name PatternType] -> [Map Name PatternType])
-> [Map Name PatternType] -> [Map Name PatternType]
forall a b. (a -> b) -> a -> b
$ (FieldBase Info VName -> Map Name PatternType)
-> [FieldBase Info VName] -> [Map Name PatternType]
forall a b. (a -> b) -> [a] -> [b]
map FieldBase Info VName -> Map Name PatternType
record [FieldBase Info VName]
fs
  where
    record :: FieldBase Info VName -> Map Name PatternType
record (RecordFieldExplicit Name
name ExpBase Info VName
e SrcLoc
_) = Name -> PatternType -> Map Name PatternType
forall k a. k -> a -> Map k a
M.singleton Name
name (PatternType -> Map Name PatternType)
-> PatternType -> Map Name PatternType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
    record (RecordFieldImplicit VName
name (Info PatternType
t) SrcLoc
_) =
      Name -> PatternType -> Map Name PatternType
forall k a. k -> a -> Map k a
M.singleton (VName -> Name
baseName VName
name) (PatternType -> Map Name PatternType)
-> PatternType -> Map Name PatternType
forall a b. (a -> b) -> a -> b
$
        PatternType
t
          PatternType -> (Aliasing -> Aliasing) -> PatternType
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` Alias -> Aliasing -> Aliasing
forall a. Ord a => a -> Set a -> Set a
S.insert (VName -> Alias
AliasBound VName
name)
typeOf (ArrayLit [ExpBase Info VName]
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (StringLit [Word8]
vs SrcLoc
_) =
  Aliasing
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> PatternType
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array
    Aliasing
forall a. Monoid a => a
mempty
    Uniqueness
Unique
    (PrimType -> ScalarTypeBase (DimDecl VName) ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (IntType -> PrimType
Unsigned IntType
Int8))
    ([DimDecl VName] -> ShapeDecl (DimDecl VName)
forall dim. [dim] -> ShapeDecl dim
ShapeDecl [Int -> DimDecl VName
forall vn. Int -> DimDecl vn
ConstDim (Int -> DimDecl VName) -> Int -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ [Word8] -> Int
forall i a. Num i => [a] -> i
genericLength [Word8]
vs])
typeOf (Range ExpBase Info VName
_ Maybe (ExpBase Info VName)
_ Inclusiveness (ExpBase Info VName)
_ (Info PatternType
t, Info [VName]
_) SrcLoc
_) = PatternType
t
typeOf (BinOp (QualName VName, SrcLoc)
_ Info PatternType
_ (ExpBase Info VName,
 Info (TypeBase (DimDecl VName) (), Maybe VName))
_ (ExpBase Info VName,
 Info (TypeBase (DimDecl VName) (), Maybe VName))
_ (Info PatternType
t) Info [VName]
_ SrcLoc
_) = PatternType
t
typeOf (Project Name
_ ExpBase Info VName
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (If ExpBase Info VName
_ ExpBase Info VName
_ ExpBase Info VName
_ (Info PatternType
t, Info [VName]
_) SrcLoc
_) = PatternType
t
typeOf (Var QualName VName
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (Ascript ExpBase Info VName
e TypeDeclBase Info VName
_ SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
typeOf (Coerce ExpBase Info VName
_ TypeDeclBase Info VName
_ (Info PatternType
t, Info [VName]
_) SrcLoc
_) = PatternType
t
typeOf (Apply ExpBase Info VName
_ ExpBase Info VName
_ Info (Diet, Maybe VName)
_ (Info PatternType
t, Info [VName]
_) SrcLoc
_) = PatternType
t
typeOf (Negate ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
typeOf (LetPat PatternBase Info VName
_ ExpBase Info VName
_ ExpBase Info VName
_ (Info PatternType
t, Info [VName]
_) SrcLoc
_) = PatternType
t
typeOf (LetFun VName
_ ([TypeParamBase VName], [PatternBase Info VName],
 Maybe (TypeExp VName), Info (TypeBase (DimDecl VName) ()),
 ExpBase Info VName)
_ ExpBase Info VName
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (LetWith IdentBase Info VName
_ IdentBase Info VName
_ [DimIndexBase Info VName]
_ ExpBase Info VName
_ ExpBase Info VName
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (Index ExpBase Info VName
_ [DimIndexBase Info VName]
_ (Info PatternType
t, Info [VName]
_) SrcLoc
_) = PatternType
t
typeOf (Update ExpBase Info VName
e [DimIndexBase Info VName]
_ ExpBase Info VName
_ SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e PatternType -> Aliasing -> PatternType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliasing
forall a. Monoid a => a
mempty
typeOf (RecordUpdate ExpBase Info VName
_ [Name]
_ ExpBase Info VName
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (Assert ExpBase Info VName
_ ExpBase Info VName
e Info String
_ SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
typeOf (DoLoop [VName]
_ PatternBase Info VName
_ ExpBase Info VName
_ LoopFormBase Info VName
_ ExpBase Info VName
_ (Info (PatternType
t, [VName]
_)) SrcLoc
_) = PatternType
t
typeOf (Lambda [PatternBase Info VName]
params ExpBase Info VName
_ Maybe (TypeExp VName)
_ (Info (Aliasing
als, TypeBase (DimDecl VName) ()
t)) SrcLoc
_) =
  Set VName -> PatternType -> PatternType
unscopeType Set VName
bound_here (PatternType -> PatternType) -> PatternType -> PatternType
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName
 -> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> [PatternBase Info VName]
-> TypeBase (DimDecl VName) ()
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((PName, TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim.
(PName, TypeBase dim ()) -> TypeBase dim () -> TypeBase dim ()
arrow ((PName, TypeBase (DimDecl VName) ())
 -> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> (PatternBase Info VName -> (PName, TypeBase (DimDecl VName) ()))
-> PatternBase Info VName
-> TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternBase Info VName -> (PName, TypeBase (DimDecl VName) ())
patternParam) TypeBase (DimDecl VName) ()
t [PatternBase Info VName]
params TypeBase (DimDecl VName) () -> Aliasing -> PatternType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliasing
als
  where
    bound_here :: Set VName
bound_here =
      (IdentBase Info VName -> VName)
-> Set (IdentBase Info VName) -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName ([Set (IdentBase Info VName)] -> Set (IdentBase Info VName)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase Info VName)] -> Set (IdentBase Info VName))
-> [Set (IdentBase Info VName)] -> Set (IdentBase Info VName)
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> Set (IdentBase Info VName))
-> [PatternBase Info VName] -> [Set (IdentBase Info VName)]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase Info VName]
params)
        Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((PatternBase Info VName -> Maybe VName)
-> [PatternBase Info VName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((PName, TypeBase (DimDecl VName) ()) -> Maybe VName
forall b. (PName, b) -> Maybe VName
named ((PName, TypeBase (DimDecl VName) ()) -> Maybe VName)
-> (PatternBase Info VName -> (PName, TypeBase (DimDecl VName) ()))
-> PatternBase Info VName
-> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternBase Info VName -> (PName, TypeBase (DimDecl VName) ())
patternParam) [PatternBase Info VName]
params)
    arrow :: (PName, TypeBase dim ()) -> TypeBase dim () -> TypeBase dim ()
arrow (PName
px, TypeBase dim ()
tx) TypeBase dim ()
y = ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> TypeBase dim ()
-> TypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
px TypeBase dim ()
tx TypeBase dim ()
y
    named :: (PName, b) -> Maybe VName
named (Named VName
x, b
_) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
x
    named (PName
Unnamed, b
_) = Maybe VName
forall a. Maybe a
Nothing
typeOf (OpSection QualName VName
_ (Info PatternType
t) SrcLoc
_) =
  PatternType
t
typeOf (OpSectionLeft QualName VName
_ Info PatternType
_ ExpBase Info VName
_ (Info (TypeBase (DimDecl VName) (), Maybe VName)
_, Info TypeBase (DimDecl VName) ()
pt2) (Info PatternType
ret, Info [VName]
_) SrcLoc
_) =
  [PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
pt2] PatternType
ret
typeOf (OpSectionRight QualName VName
_ Info PatternType
_ ExpBase Info VName
_ (Info TypeBase (DimDecl VName) ()
pt1, Info (TypeBase (DimDecl VName) (), Maybe VName)
_) (Info PatternType
ret) SrcLoc
_) =
  [PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
pt1] PatternType
ret
typeOf (ProjectSection [Name]
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (IndexSection [DimIndexBase Info VName]
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (Constr Name
_ [ExpBase Info VName]
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
typeOf (Match ExpBase Info VName
_ NonEmpty (CaseBase Info VName)
cs (Info PatternType
t, Info [VName]
_) SrcLoc
_) =
  Set VName -> PatternType -> PatternType
unscopeType ((CaseBase Info VName -> Set VName)
-> NonEmpty (CaseBase Info VName) -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap CaseBase Info VName -> Set VName
forall b (f :: * -> *). (Ord b, Functor f) => CaseBase f b -> Set b
unscopeSet NonEmpty (CaseBase Info VName)
cs) PatternType
t
  where
    unscopeSet :: CaseBase f b -> Set b
unscopeSet (CasePat PatternBase f b
p ExpBase f b
_ SrcLoc
_) = (IdentBase f b -> b) -> Set (IdentBase f b) -> Set b
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase f b -> b
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (Set (IdentBase f b) -> Set b) -> Set (IdentBase f b) -> Set b
forall a b. (a -> b) -> a -> b
$ PatternBase f b -> Set (IdentBase f b)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents PatternBase f b
p
typeOf (Attr AttrInfo
_ ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e

-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes
-- @ts@ as parameters and returns @ret@.
foldFunType :: Monoid as => [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType :: [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase dim as]
ps TypeBase dim as
ret = (TypeBase dim as -> TypeBase dim as -> TypeBase dim as)
-> TypeBase dim as -> [TypeBase dim as] -> TypeBase dim as
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TypeBase dim as -> TypeBase dim as -> TypeBase dim as
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
arrow TypeBase dim as
ret [TypeBase dim as]
ps
  where
    arrow :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
arrow TypeBase dim as
t1 TypeBase dim as
t2 = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty PName
Unnamed TypeBase dim as
t1 TypeBase dim as
t2

-- | Extract the parameter types and return type from a type.
-- If the type is not an arrow type, the list of parameter types is empty.
unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType (Scalar (Arrow as
_ PName
_ TypeBase dim as
t1 TypeBase dim as
t2)) =
  let ([TypeBase dim as]
ps, TypeBase dim as
r) = TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
forall dim as.
TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType TypeBase dim as
t2
   in (TypeBase dim as
t1 TypeBase dim as -> [TypeBase dim as] -> [TypeBase dim as]
forall a. a -> [a] -> [a]
: [TypeBase dim as]
ps, TypeBase dim as
r)
unfoldFunType TypeBase dim as
t = ([], TypeBase dim as
t)

-- | The type names mentioned in a type.
typeVars :: Monoid as => TypeBase dim as -> S.Set VName
typeVars :: TypeBase dim as -> Set VName
typeVars TypeBase dim as
t =
  case TypeBase dim as
t of
    Scalar Prim {} -> Set VName
forall a. Monoid a => a
mempty
    Scalar (TypeVar as
_ Uniqueness
_ TypeName
tn [TypeArg dim]
targs) ->
      [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ TypeName -> Set VName
typeVarFree TypeName
tn Set VName -> [Set VName] -> [Set VName]
forall a. a -> [a] -> [a]
: (TypeArg dim -> Set VName) -> [TypeArg dim] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg dim -> Set VName
forall dim. TypeArg dim -> Set VName
typeArgFree [TypeArg dim]
targs
    Scalar (Arrow as
_ PName
_ TypeBase dim as
t1 TypeBase dim as
t2) -> TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim as
t1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim as
t2
    Scalar (Record Map Name (TypeBase dim as)
fields) -> (TypeBase dim as -> Set VName)
-> Map Name (TypeBase dim as) -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars Map Name (TypeBase dim as)
fields
    Scalar (Sum Map Name [TypeBase dim as]
cs) -> [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (([TypeBase dim as] -> [Set VName])
-> Map Name [TypeBase dim as] -> [Set VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (([TypeBase dim as] -> [Set VName])
 -> Map Name [TypeBase dim as] -> [Set VName])
-> ((TypeBase dim as -> Set VName)
    -> [TypeBase dim as] -> [Set VName])
-> (TypeBase dim as -> Set VName)
-> Map Name [TypeBase dim as]
-> [Set VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase dim as -> Set VName) -> [TypeBase dim as] -> [Set VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars Map Name [TypeBase dim as]
cs
    Array as
_ Uniqueness
_ ScalarTypeBase dim ()
rt ShapeDecl dim
_ -> TypeBase dim () -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars (TypeBase dim () -> Set VName) -> TypeBase dim () -> Set VName
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
rt
  where
    typeVarFree :: TypeName -> Set VName
typeVarFree = VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName)
-> (TypeName -> VName) -> TypeName -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeName -> VName
typeLeaf
    typeArgFree :: TypeArg dim -> Set VName
typeArgFree (TypeArgType TypeBase dim ()
ta SrcLoc
_) = TypeBase dim () -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim ()
ta
    typeArgFree TypeArgDim {} = Set VName
forall a. Monoid a => a
mempty

-- | @orderZero t@ is 'True' if the argument type has order 0, i.e., it is not
-- a function type, does not contain a function type as a subcomponent, and may
-- not be instantiated with a function type.
orderZero :: TypeBase dim as -> Bool
orderZero :: TypeBase dim as -> Bool
orderZero Array {} = Bool
True
orderZero (Scalar (Prim PrimType
_)) = Bool
True
orderZero (Scalar (Record Map Name (TypeBase dim as)
fs)) = (TypeBase dim as -> Bool) -> [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim as -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero ([TypeBase dim as] -> Bool) -> [TypeBase dim as] -> Bool
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> [TypeBase dim as]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase dim as)
fs
orderZero (Scalar TypeVar {}) = Bool
True
orderZero (Scalar Arrow {}) = Bool
False
orderZero (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
all ((TypeBase dim as -> Bool) -> [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim as -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero) Map Name [TypeBase dim as]
cs

-- | Extract all the shape names that occur in a given pattern.
patternDimNames :: PatternBase Info VName -> S.Set VName
patternDimNames :: PatternBase Info VName -> Set VName
patternDimNames (TuplePattern [PatternBase Info VName]
ps SrcLoc
_) = (PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
patternDimNames [PatternBase Info VName]
ps
patternDimNames (RecordPattern [(Name, PatternBase Info VName)]
fs SrcLoc
_) = ((Name, PatternBase Info VName) -> Set VName)
-> [(Name, PatternBase Info VName)] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatternBase Info VName -> Set VName
patternDimNames (PatternBase Info VName -> Set VName)
-> ((Name, PatternBase Info VName) -> PatternBase Info VName)
-> (Name, PatternBase Info VName)
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatternBase Info VName) -> PatternBase Info VName
forall a b. (a, b) -> b
snd) [(Name, PatternBase Info VName)]
fs
patternDimNames (PatternParens PatternBase Info VName
p SrcLoc
_) = PatternBase Info VName -> Set VName
patternDimNames PatternBase Info VName
p
patternDimNames (Id VName
_ (Info PatternType
tp) SrcLoc
_) = PatternType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames PatternType
tp
patternDimNames (Wildcard (Info PatternType
tp) SrcLoc
_) = PatternType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames PatternType
tp
patternDimNames (PatternAscription PatternBase Info VName
p (TypeDecl TypeExp VName
_ (Info TypeBase (DimDecl VName) ()
t)) SrcLoc
_) =
  PatternBase Info VName -> Set VName
patternDimNames PatternBase Info VName
p Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> TypeBase (DimDecl VName) () -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames TypeBase (DimDecl VName) ()
t
patternDimNames (PatternLit PatLit
_ (Info PatternType
tp) SrcLoc
_) = PatternType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames PatternType
tp
patternDimNames (PatternConstr Name
_ Info PatternType
_ [PatternBase Info VName]
ps SrcLoc
_) = (PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
patternDimNames [PatternBase Info VName]
ps

-- | Extract all the shape names that occur in a given type.
typeDimNames :: TypeBase (DimDecl VName) als -> S.Set VName
typeDimNames :: TypeBase (DimDecl VName) als -> Set VName
typeDimNames = (DimDecl VName -> Set VName) -> [DimDecl VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimDecl VName -> Set VName
dimName ([DimDecl VName] -> Set VName)
-> (TypeBase (DimDecl VName) als -> [DimDecl VName])
-> TypeBase (DimDecl VName) als
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) als -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims
  where
    dimName :: DimDecl VName -> S.Set VName
    dimName :: DimDecl VName -> Set VName
dimName (NamedDim QualName VName
qn) = VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
    dimName DimDecl VName
_ = Set VName
forall a. Monoid a => a
mempty

-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern
-- have order 0.
patternOrderZero :: PatternBase Info vn -> Bool
patternOrderZero :: PatternBase Info vn -> Bool
patternOrderZero PatternBase Info vn
pat = case PatternBase Info vn
pat of
  TuplePattern [PatternBase Info vn]
ps SrcLoc
_ -> (PatternBase Info vn -> Bool) -> [PatternBase Info vn] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all PatternBase Info vn -> Bool
forall vn. PatternBase Info vn -> Bool
patternOrderZero [PatternBase Info vn]
ps
  RecordPattern [(Name, PatternBase Info vn)]
fs SrcLoc
_ -> ((Name, PatternBase Info vn) -> Bool)
-> [(Name, PatternBase Info vn)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (PatternBase Info vn -> Bool
forall vn. PatternBase Info vn -> Bool
patternOrderZero (PatternBase Info vn -> Bool)
-> ((Name, PatternBase Info vn) -> PatternBase Info vn)
-> (Name, PatternBase Info vn)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatternBase Info vn) -> PatternBase Info vn
forall a b. (a, b) -> b
snd) [(Name, PatternBase Info vn)]
fs
  PatternParens PatternBase Info vn
p SrcLoc
_ -> PatternBase Info vn -> Bool
forall vn. PatternBase Info vn -> Bool
patternOrderZero PatternBase Info vn
p
  Id vn
_ (Info PatternType
t) SrcLoc
_ -> PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
t
  Wildcard (Info PatternType
t) SrcLoc
_ -> PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
t
  PatternAscription PatternBase Info vn
p TypeDeclBase Info vn
_ SrcLoc
_ -> PatternBase Info vn -> Bool
forall vn. PatternBase Info vn -> Bool
patternOrderZero PatternBase Info vn
p
  PatternLit PatLit
_ (Info PatternType
t) SrcLoc
_ -> PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
t
  PatternConstr Name
_ Info PatternType
_ [PatternBase Info vn]
ps SrcLoc
_ -> (PatternBase Info vn -> Bool) -> [PatternBase Info vn] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all PatternBase Info vn -> Bool
forall vn. PatternBase Info vn -> Bool
patternOrderZero [PatternBase Info vn]
ps

-- | The set of identifiers bound in a pattern.
patternIdents :: (Functor f, Ord vn) => PatternBase f vn -> S.Set (IdentBase f vn)
patternIdents :: PatternBase f vn -> Set (IdentBase f vn)
patternIdents (Id vn
v f PatternType
t SrcLoc
loc) = IdentBase f vn -> Set (IdentBase f vn)
forall a. a -> Set a
S.singleton (IdentBase f vn -> Set (IdentBase f vn))
-> IdentBase f vn -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ vn -> f PatternType -> SrcLoc -> IdentBase f vn
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> IdentBase f vn
Ident vn
v f PatternType
t SrcLoc
loc
patternIdents (PatternParens PatternBase f vn
p SrcLoc
_) = PatternBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents PatternBase f vn
p
patternIdents (TuplePattern [PatternBase f vn]
pats SrcLoc
_) = [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase f vn)] -> Set (IdentBase f vn))
-> [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ (PatternBase f vn -> Set (IdentBase f vn))
-> [PatternBase f vn] -> [Set (IdentBase f vn)]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase f vn]
pats
patternIdents (RecordPattern [(Name, PatternBase f vn)]
fs SrcLoc
_) = [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase f vn)] -> Set (IdentBase f vn))
-> [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ ((Name, PatternBase f vn) -> Set (IdentBase f vn))
-> [(Name, PatternBase f vn)] -> [Set (IdentBase f vn)]
forall a b. (a -> b) -> [a] -> [b]
map (PatternBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents (PatternBase f vn -> Set (IdentBase f vn))
-> ((Name, PatternBase f vn) -> PatternBase f vn)
-> (Name, PatternBase f vn)
-> Set (IdentBase f vn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatternBase f vn) -> PatternBase f vn
forall a b. (a, b) -> b
snd) [(Name, PatternBase f vn)]
fs
patternIdents Wildcard {} = Set (IdentBase f vn)
forall a. Monoid a => a
mempty
patternIdents (PatternAscription PatternBase f vn
p TypeDeclBase f vn
_ SrcLoc
_) = PatternBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents PatternBase f vn
p
patternIdents PatternLit {} = Set (IdentBase f vn)
forall a. Monoid a => a
mempty
patternIdents (PatternConstr Name
_ f PatternType
_ [PatternBase f vn]
ps SrcLoc
_) = [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase f vn)] -> Set (IdentBase f vn))
-> [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ (PatternBase f vn -> Set (IdentBase f vn))
-> [PatternBase f vn] -> [Set (IdentBase f vn)]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase f vn]
ps

-- | The set of names bound in a pattern.
patternNames :: (Functor f, Ord vn) => PatternBase f vn -> S.Set vn
patternNames :: PatternBase f vn -> Set vn
patternNames (Id vn
v f PatternType
_ SrcLoc
_) = vn -> Set vn
forall a. a -> Set a
S.singleton vn
v
patternNames (PatternParens PatternBase f vn
p SrcLoc
_) = PatternBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames PatternBase f vn
p
patternNames (TuplePattern [PatternBase f vn]
pats SrcLoc
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ (PatternBase f vn -> Set vn) -> [PatternBase f vn] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames [PatternBase f vn]
pats
patternNames (RecordPattern [(Name, PatternBase f vn)]
fs SrcLoc
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ ((Name, PatternBase f vn) -> Set vn)
-> [(Name, PatternBase f vn)] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map (PatternBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames (PatternBase f vn -> Set vn)
-> ((Name, PatternBase f vn) -> PatternBase f vn)
-> (Name, PatternBase f vn)
-> Set vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatternBase f vn) -> PatternBase f vn
forall a b. (a, b) -> b
snd) [(Name, PatternBase f vn)]
fs
patternNames Wildcard {} = Set vn
forall a. Monoid a => a
mempty
patternNames (PatternAscription PatternBase f vn
p TypeDeclBase f vn
_ SrcLoc
_) = PatternBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames PatternBase f vn
p
patternNames PatternLit {} = Set vn
forall a. Monoid a => a
mempty
patternNames (PatternConstr Name
_ f PatternType
_ [PatternBase f vn]
ps SrcLoc
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ (PatternBase f vn -> Set vn) -> [PatternBase f vn] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames [PatternBase f vn]
ps

-- | A mapping from names bound in a map to their identifier.
patternMap :: (Functor f) => PatternBase f VName -> M.Map VName (IdentBase f VName)
patternMap :: PatternBase f VName -> Map VName (IdentBase f VName)
patternMap PatternBase f VName
pat =
  [(VName, IdentBase f VName)] -> Map VName (IdentBase f VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, IdentBase f VName)] -> Map VName (IdentBase f VName))
-> [(VName, IdentBase f VName)] -> Map VName (IdentBase f VName)
forall a b. (a -> b) -> a -> b
$ [VName] -> [IdentBase f VName] -> [(VName, IdentBase f VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((IdentBase f VName -> VName) -> [IdentBase f VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map IdentBase f VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName [IdentBase f VName]
idents) [IdentBase f VName]
idents
  where
    idents :: [IdentBase f VName]
idents = Set (IdentBase f VName) -> [IdentBase f VName]
forall a. Set a -> [a]
S.toList (Set (IdentBase f VName) -> [IdentBase f VName])
-> Set (IdentBase f VName) -> [IdentBase f VName]
forall a b. (a -> b) -> a -> b
$ PatternBase f VName -> Set (IdentBase f VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents PatternBase f VName
pat

-- | The type of values bound by the pattern.
patternType :: PatternBase Info VName -> PatternType
patternType :: PatternBase Info VName -> PatternType
patternType (Wildcard (Info PatternType
t) SrcLoc
_) = PatternType
t
patternType (PatternParens PatternBase Info VName
p SrcLoc
_) = PatternBase Info VName -> PatternType
patternType PatternBase Info VName
p
patternType (Id VName
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
patternType (TuplePattern [PatternBase Info VName]
pats SrcLoc
_) = [PatternType] -> PatternType
forall dim as. [TypeBase dim as] -> TypeBase dim as
tupleRecord ([PatternType] -> PatternType) -> [PatternType] -> PatternType
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> PatternType)
-> [PatternBase Info VName] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase Info VName -> PatternType
patternType [PatternBase Info VName]
pats
patternType (RecordPattern [(Name, PatternBase Info VName)]
fs SrcLoc
_) = ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ Map Name PatternType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name PatternType -> ScalarTypeBase (DimDecl VName) Aliasing)
-> Map Name PatternType -> ScalarTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ PatternBase Info VName -> PatternType
patternType (PatternBase Info VName -> PatternType)
-> Map Name (PatternBase Info VName) -> Map Name PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Name, PatternBase Info VName)]
-> Map Name (PatternBase Info VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatternBase Info VName)]
fs
patternType (PatternAscription PatternBase Info VName
p TypeDeclBase Info VName
_ SrcLoc
_) = PatternBase Info VName -> PatternType
patternType PatternBase Info VName
p
patternType (PatternLit PatLit
_ (Info PatternType
t) SrcLoc
_) = PatternType
t
patternType (PatternConstr Name
_ (Info PatternType
t) [PatternBase Info VName]
_ SrcLoc
_) = PatternType
t

-- | The type matched by the pattern, including shape declarations if present.
patternStructType :: PatternBase Info VName -> StructType
patternStructType :: PatternBase Info VName -> TypeBase (DimDecl VName) ()
patternStructType = PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatternType -> TypeBase (DimDecl VName) ())
-> (PatternBase Info VName -> PatternType)
-> PatternBase Info VName
-> TypeBase (DimDecl VName) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternBase Info VName -> PatternType
patternType

-- | When viewed as a function parameter, does this pattern correspond
-- to a named parameter of some type?
patternParam :: PatternBase Info VName -> (PName, StructType)
patternParam :: PatternBase Info VName -> (PName, TypeBase (DimDecl VName) ())
patternParam (PatternParens PatternBase Info VName
p SrcLoc
_) =
  PatternBase Info VName -> (PName, TypeBase (DimDecl VName) ())
patternParam PatternBase Info VName
p
patternParam (PatternAscription (Id VName
v Info PatternType
_ SrcLoc
_) TypeDeclBase Info VName
td SrcLoc
_) =
  (VName -> PName
Named VName
v, Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ()
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info (TypeBase (DimDecl VName) ())
forall (f :: * -> *) vn.
TypeDeclBase f vn -> f (TypeBase (DimDecl VName) ())
expandedType TypeDeclBase Info VName
td)
patternParam (Id VName
v (Info PatternType
t) SrcLoc
_) =
  (VName -> PName
Named VName
v, PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t)
patternParam PatternBase Info VName
p =
  (PName
Unnamed, PatternBase Info VName -> TypeBase (DimDecl VName) ()
patternStructType PatternBase Info VName
p)

-- | Names of primitive types to types.  This is only valid if no
-- shadowing is going on, but useful for tools.
namesToPrimTypes :: M.Map Name PrimType
namesToPrimTypes :: Map Name PrimType
namesToPrimTypes =
  [(Name, PrimType)] -> Map Name PrimType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    [ (String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ PrimType -> String
forall a. Pretty a => a -> String
pretty PrimType
t, PrimType
t)
      | PrimType
t <-
          PrimType
Bool PrimType -> [PrimType] -> [PrimType]
forall a. a -> [a] -> [a]
:
          (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
            [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
            [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (FloatType -> PrimType) -> [FloatType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [FloatType
forall a. Bounded a => a
minBound .. FloatType
forall a. Bounded a => a
maxBound]
    ]

-- | The nature of something predefined.  These can either be
-- monomorphic or overloaded.  An overloaded builtin is a list valid
-- types it can be instantiated with, to the parameter and result
-- type, with 'Nothing' representing the overloaded parameter type.
data Intrinsic
  = IntrinsicMonoFun [PrimType] PrimType
  | IntrinsicOverloadedFun [PrimType] [Maybe PrimType] (Maybe PrimType)
  | IntrinsicPolyFun [TypeParamBase VName] [StructType] StructType
  | IntrinsicType PrimType
  | IntrinsicEquality -- Special cased.

-- | A map of all built-ins.
intrinsics :: M.Map VName Intrinsic
intrinsics :: Map VName Intrinsic
intrinsics =
  [(VName, Intrinsic)] -> Map VName Intrinsic
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Intrinsic)] -> Map VName Intrinsic)
-> [(VName, Intrinsic)] -> Map VName Intrinsic
forall a b. (a -> b) -> a -> b
$
    (Int -> (String, Intrinsic) -> (VName, Intrinsic))
-> [Int] -> [(String, Intrinsic)] -> [(VName, Intrinsic)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> (String, Intrinsic) -> (VName, Intrinsic)
forall b. Int -> (String, b) -> (VName, b)
namify [Int
10 ..] ([(String, Intrinsic)] -> [(VName, Intrinsic)])
-> [(String, Intrinsic)] -> [(VName, Intrinsic)]
forall a b. (a -> b) -> a -> b
$
      ((String, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
 -> (String, Intrinsic))
-> [(String,
     ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
-> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map (String, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
-> (String, Intrinsic)
forall a c. (a, ([PrimType], PrimType, c)) -> (a, Intrinsic)
primFun (Map String ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> [(String,
     ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
forall k a. Map k a -> [(k, a)]
M.toList Map String ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
Primitive.primFuns)
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ [(String
"opaque", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a] [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a] (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_a)]
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (UnOp -> (String, Intrinsic)) -> [UnOp] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map UnOp -> (String, Intrinsic)
unOpFun [UnOp]
Primitive.allUnOps
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (BinOp -> (String, Intrinsic)) -> [BinOp] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map BinOp -> (String, Intrinsic)
binOpFun [BinOp]
Primitive.allBinOps
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (CmpOp -> (String, Intrinsic)) -> [CmpOp] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map CmpOp -> (String, Intrinsic)
cmpOpFun [CmpOp]
Primitive.allCmpOps
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (ConvOp -> (String, Intrinsic))
-> [ConvOp] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map ConvOp -> (String, Intrinsic)
convOpFun [ConvOp]
Primitive.allConvOps
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (IntType -> (String, Intrinsic))
-> [IntType] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> (String, Intrinsic)
signFun [IntType]
Primitive.allIntTypes
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (IntType -> (String, Intrinsic))
-> [IntType] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> (String, Intrinsic)
unsignFun [IntType]
Primitive.allIntTypes
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (PrimType -> (String, Intrinsic))
-> [PrimType] -> [(String, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map
          PrimType -> (String, Intrinsic)
intrinsicType
          ( (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
              [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
              [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (FloatType -> PrimType) -> [FloatType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [FloatType
forall a. Bounded a => a
minBound .. FloatType
forall a. Bounded a => a
maxBound]
              [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ [PrimType
Bool]
          )
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++
        -- This overrides the ! from Primitive.
        [ ( String
"!",
            [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> Intrinsic
IntrinsicOverloadedFun
              ( (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
                  [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
                  [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ [PrimType
Bool]
              )
              [Maybe PrimType
forall a. Maybe a
Nothing]
              Maybe PrimType
forall a. Maybe a
Nothing
          )
        ]
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++
        -- The reason for the loop formulation is to ensure that we
        -- get a missing case warning if we forget a case.
        (BinOp -> Maybe (String, Intrinsic))
-> [BinOp] -> [(String, Intrinsic)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe BinOp -> Maybe (String, Intrinsic)
mkIntrinsicBinOp [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound]
        [(String, Intrinsic)]
-> [(String, Intrinsic)] -> [(String, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ [ ( String
"flatten",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
2)]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
forall a b. (a -> b) -> a -> b
$ ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
             ),
             ( String
"unflatten",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ 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,
                   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,
                   ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
                 ]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
forall a b. (a -> b) -> a -> b
$ ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
2)
             ),
             ( String
"concat",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [TypeBase (DimDecl VName) ()
arr_a, TypeBase (DimDecl VName) ()
arr_a]
                 TypeBase (DimDecl VName) ()
uarr_a
             ),
             ( String
"rotate",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [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, TypeBase (DimDecl VName) ()
arr_a]
                 TypeBase (DimDecl VName) ()
arr_a
             ),
             (String
"transpose", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a] [TypeBase (DimDecl VName) ()
arr_2d_a] TypeBase (DimDecl VName) ()
arr_2d_a),
             ( String
"scatter",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Unique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1),
                   ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique (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) (Int -> ShapeDecl (DimDecl VName)
rank Int
1),
                   ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
                 ]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
forall a b. (a -> b) -> a -> b
$ ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Unique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
             ),
             (String
"zip", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b] [TypeBase (DimDecl VName) ()
arr_a, TypeBase (DimDecl VName) ()
arr_b] TypeBase (DimDecl VName) ()
arr_a_b),
             (String
"unzip", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b] [TypeBase (DimDecl VName) ()
arr_a_b] TypeBase (DimDecl VName) ()
t_arr_a_arr_b),
             ( String
"hist",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ 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,
                   TypeBase (DimDecl VName) ()
uarr_a,
                   ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a),
                   ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a,
                   ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique (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) (Int -> ShapeDecl (DimDecl VName)
rank Int
1),
                   TypeBase (DimDecl VName) ()
arr_a
                 ]
                 TypeBase (DimDecl VName) ()
uarr_a
             ),
             (String
"map", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b] [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b, TypeBase (DimDecl VName) ()
arr_a] TypeBase (DimDecl VName) ()
uarr_b),
             ( String
"reduce",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a), ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a, TypeBase (DimDecl VName) ()
arr_a]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_a
             ),
             ( String
"reduce_comm",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a), ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a, TypeBase (DimDecl VName) ()
arr_a]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_a
             ),
             ( String
"scan",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a), ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a, TypeBase (DimDecl VName) ()
arr_a]
                 TypeBase (DimDecl VName) ()
uarr_a
             ),
             ( String
"partition",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a]
                 [ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (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
Int32),
                   ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (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),
                   TypeBase (DimDecl VName) ()
arr_a
                 ]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [TypeBase (DimDecl VName) ()] -> TypeBase (DimDecl VName) ()
forall dim as. [TypeBase dim as] -> TypeBase dim as
tupleRecord [TypeBase (DimDecl VName) ()
uarr_a, ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Unique (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) (Int -> ShapeDecl (DimDecl VName)
rank Int
1)]
             ),
             ( String
"map_stream",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                 [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (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) TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`karr` (TypeBase (DimDecl VName) ()
arr_ka TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` TypeBase (DimDecl VName) ()
arr_kb), TypeBase (DimDecl VName) ()
arr_a]
                 TypeBase (DimDecl VName) ()
uarr_b
             ),
             ( String
"map_stream_per",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                 [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (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) TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`karr` (TypeBase (DimDecl VName) ()
arr_ka TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` TypeBase (DimDecl VName) ()
arr_kb), TypeBase (DimDecl VName) ()
arr_a]
                 TypeBase (DimDecl VName) ()
uarr_b
             ),
             ( String
"reduce_stream",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                 [ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b),
                   ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (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) TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`karr` (TypeBase (DimDecl VName) ()
arr_ka TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b),
                   TypeBase (DimDecl VName) ()
arr_a
                 ]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_b
             ),
             ( String
"reduce_stream_per",
               [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun
                 [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                 [ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` (ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b),
                   ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (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) TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`karr` (TypeBase (DimDecl VName) ()
arr_ka TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
`arr` ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b),
                   TypeBase (DimDecl VName) ()
arr_a
                 ]
                 (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_b
             ),
             (String
"trace", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a] [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a] (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_a),
             (String
"break", [TypeParamBase VName]
-> [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) ()
-> Intrinsic
IntrinsicPolyFun [TypeParamBase VName
tp_a] [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a] (TypeBase (DimDecl VName) () -> Intrinsic)
-> TypeBase (DimDecl VName) () -> Intrinsic
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) ()
forall dim. ScalarTypeBase dim ()
t_a)
           ]
  where
    tv_a :: VName
tv_a = Name -> Int -> VName
VName (String -> Name
nameFromString String
"a") Int
0
    t_a :: ScalarTypeBase dim ()
t_a = ()
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim ()
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> TypeName
typeName VName
tv_a) []
    arr_a :: TypeBase (DimDecl VName) ()
arr_a = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
    arr_2d_a :: TypeBase (DimDecl VName) ()
arr_2d_a = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
2)
    uarr_a :: TypeBase (DimDecl VName) ()
uarr_a = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Unique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
    tp_a :: TypeParamBase VName
tp_a = Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
tv_a SrcLoc
forall a. Monoid a => a
mempty

    tv_b :: VName
tv_b = Name -> Int -> VName
VName (String -> Name
nameFromString String
"b") Int
1
    t_b :: ScalarTypeBase dim ()
t_b = ()
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim ()
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> TypeName
typeName VName
tv_b) []
    arr_b :: TypeBase (DimDecl VName) ()
arr_b = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
    uarr_b :: TypeBase (DimDecl VName) ()
uarr_b = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Unique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
    tp_b :: TypeParamBase VName
tp_b = Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
tv_b SrcLoc
forall a. Monoid a => a
mempty

    arr_a_b :: TypeBase (DimDecl VName) ()
arr_a_b =
      ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array
        ()
        Uniqueness
Nonunique
        (Map Name (TypeBase (DimDecl VName) ())
-> ScalarTypeBase (DimDecl VName) ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record ([(Name, TypeBase (DimDecl VName) ())]
-> Map Name (TypeBase (DimDecl VName) ())
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase (DimDecl VName) ())]
 -> Map Name (TypeBase (DimDecl VName) ()))
-> [(Name, TypeBase (DimDecl VName) ())]
-> Map Name (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$ [Name]
-> [TypeBase (DimDecl VName) ()]
-> [(Name, TypeBase (DimDecl VName) ())]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a, ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b]))
        (Int -> ShapeDecl (DimDecl VName)
rank Int
1)
    t_arr_a_arr_b :: TypeBase (DimDecl VName) ()
t_arr_a_arr_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
$ [(Name, TypeBase (DimDecl VName) ())]
-> Map Name (TypeBase (DimDecl VName) ())
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase (DimDecl VName) ())]
 -> Map Name (TypeBase (DimDecl VName) ()))
-> [(Name, TypeBase (DimDecl VName) ())]
-> Map Name (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$ [Name]
-> [TypeBase (DimDecl VName) ()]
-> [(Name, TypeBase (DimDecl VName) ())]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [TypeBase (DimDecl VName) ()
arr_a, TypeBase (DimDecl VName) ()
arr_b]

    arr :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
arr TypeBase dim as
x TypeBase dim as
y = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty PName
Unnamed TypeBase dim as
x TypeBase dim as
y

    kv :: VName
kv = Name -> Int -> VName
VName (String -> Name
nameFromString String
"k") Int
2
    arr_ka :: TypeBase (DimDecl VName) ()
arr_ka = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_a ([DimDecl VName] -> ShapeDecl (DimDecl VName)
forall dim. [dim] -> ShapeDecl dim
ShapeDecl [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
kv])
    arr_kb :: TypeBase (DimDecl VName) ()
arr_kb = ()
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
Nonunique ScalarTypeBase (DimDecl VName) ()
forall dim. ScalarTypeBase dim ()
t_b ([DimDecl VName] -> ShapeDecl (DimDecl VName)
forall dim. [dim] -> ShapeDecl dim
ShapeDecl [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
kv])
    karr :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
karr TypeBase dim as
x TypeBase dim as
y = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty (VName -> PName
Named VName
kv) TypeBase dim as
x TypeBase dim as
y

    namify :: Int -> (String, b) -> (VName, b)
namify Int
i (String
k, b
v) = (Name -> Int -> VName
VName (String -> Name
nameFromString String
k) Int
i, b
v)

    primFun :: (a, ([PrimType], PrimType, c)) -> (a, Intrinsic)
primFun (a
name, ([PrimType]
ts, PrimType
t, c
_)) =
      (a
name, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun ((PrimType -> PrimType) -> [PrimType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> PrimType
unPrim [PrimType]
ts) (PrimType -> Intrinsic) -> PrimType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
unPrim PrimType
t)

    unOpFun :: UnOp -> (String, Intrinsic)
unOpFun UnOp
bop = (UnOp -> String
forall a. Pretty a => a -> String
pretty UnOp
bop, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [PrimType
t] PrimType
t)
      where
        t :: PrimType
t = PrimType -> PrimType
unPrim (PrimType -> PrimType) -> PrimType -> PrimType
forall a b. (a -> b) -> a -> b
$ UnOp -> PrimType
Primitive.unOpType UnOp
bop

    binOpFun :: BinOp -> (String, Intrinsic)
binOpFun BinOp
bop = (BinOp -> String
forall a. Pretty a => a -> String
pretty BinOp
bop, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [PrimType
t, PrimType
t] PrimType
t)
      where
        t :: PrimType
t = PrimType -> PrimType
unPrim (PrimType -> PrimType) -> PrimType -> PrimType
forall a b. (a -> b) -> a -> b
$ BinOp -> PrimType
Primitive.binOpType BinOp
bop

    cmpOpFun :: CmpOp -> (String, Intrinsic)
cmpOpFun CmpOp
bop = (CmpOp -> String
forall a. Pretty a => a -> String
pretty CmpOp
bop, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [PrimType
t, PrimType
t] PrimType
Bool)
      where
        t :: PrimType
t = PrimType -> PrimType
unPrim (PrimType -> PrimType) -> PrimType -> PrimType
forall a b. (a -> b) -> a -> b
$ CmpOp -> PrimType
Primitive.cmpOpType CmpOp
bop

    convOpFun :: ConvOp -> (String, Intrinsic)
convOpFun ConvOp
cop = (ConvOp -> String
forall a. Pretty a => a -> String
pretty ConvOp
cop, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [PrimType -> PrimType
unPrim PrimType
ft] (PrimType -> Intrinsic) -> PrimType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
unPrim PrimType
tt)
      where
        (PrimType
ft, PrimType
tt) = ConvOp -> (PrimType, PrimType)
Primitive.convOpType ConvOp
cop

    signFun :: IntType -> (String, Intrinsic)
signFun IntType
t = (String
"sign_" String -> ShowS
forall a. [a] -> [a] -> [a]
++ IntType -> String
forall a. Pretty a => a -> String
pretty IntType
t, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [IntType -> PrimType
Unsigned IntType
t] (PrimType -> Intrinsic) -> PrimType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
t)

    unsignFun :: IntType -> (String, Intrinsic)
unsignFun IntType
t = (String
"unsign_" String -> ShowS
forall a. [a] -> [a] -> [a]
++ IntType -> String
forall a. Pretty a => a -> String
pretty IntType
t, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [IntType -> PrimType
Signed IntType
t] (PrimType -> Intrinsic) -> PrimType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
t)

    unPrim :: PrimType -> PrimType
unPrim (Primitive.IntType IntType
t) = IntType -> PrimType
Signed IntType
t
    unPrim (Primitive.FloatType FloatType
t) = FloatType -> PrimType
FloatType FloatType
t
    unPrim PrimType
Primitive.Bool = PrimType
Bool
    unPrim PrimType
Primitive.Cert = PrimType
Bool

    intrinsicType :: PrimType -> (String, Intrinsic)
intrinsicType PrimType
t = (PrimType -> String
forall a. Pretty a => a -> String
pretty PrimType
t, PrimType -> Intrinsic
IntrinsicType PrimType
t)

    anyIntType :: [PrimType]
anyIntType =
      (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
        [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
    anyNumberType :: [PrimType]
anyNumberType =
      [PrimType]
anyIntType
        [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (FloatType -> PrimType) -> [FloatType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [FloatType
forall a. Bounded a => a
minBound .. FloatType
forall a. Bounded a => a
maxBound]
    anyPrimType :: [PrimType]
anyPrimType = PrimType
Bool PrimType -> [PrimType] -> [PrimType]
forall a. a -> [a] -> [a]
: [PrimType]
anyNumberType

    mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic)
    mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic)
mkIntrinsicBinOp BinOp
op = do
      Intrinsic
op' <- BinOp -> Maybe Intrinsic
intrinsicBinOp BinOp
op
      (String, Intrinsic) -> Maybe (String, Intrinsic)
forall (m :: * -> *) a. Monad m => a -> m a
return (BinOp -> String
forall a. Pretty a => a -> String
pretty BinOp
op, Intrinsic
op')

    binOp :: [PrimType] -> Maybe Intrinsic
binOp [PrimType]
ts = Intrinsic -> Maybe Intrinsic
forall a. a -> Maybe a
Just (Intrinsic -> Maybe Intrinsic) -> Intrinsic -> Maybe Intrinsic
forall a b. (a -> b) -> a -> b
$ [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> Intrinsic
IntrinsicOverloadedFun [PrimType]
ts [Maybe PrimType
forall a. Maybe a
Nothing, Maybe PrimType
forall a. Maybe a
Nothing] Maybe PrimType
forall a. Maybe a
Nothing
    ordering :: Maybe Intrinsic
ordering = Intrinsic -> Maybe Intrinsic
forall a. a -> Maybe a
Just (Intrinsic -> Maybe Intrinsic) -> Intrinsic -> Maybe Intrinsic
forall a b. (a -> b) -> a -> b
$ [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> Intrinsic
IntrinsicOverloadedFun [PrimType]
anyPrimType [Maybe PrimType
forall a. Maybe a
Nothing, Maybe PrimType
forall a. Maybe a
Nothing] (PrimType -> Maybe PrimType
forall a. a -> Maybe a
Just PrimType
Bool)

    intrinsicBinOp :: BinOp -> Maybe Intrinsic
intrinsicBinOp BinOp
Plus = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyNumberType
    intrinsicBinOp BinOp
Minus = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyNumberType
    intrinsicBinOp BinOp
Pow = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyNumberType
    intrinsicBinOp BinOp
Times = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyNumberType
    intrinsicBinOp BinOp
Divide = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyNumberType
    intrinsicBinOp BinOp
Mod = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyNumberType
    intrinsicBinOp BinOp
Quot = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
Rem = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
ShiftR = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
ShiftL = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
Band = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
Xor = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
Bor = [PrimType] -> Maybe Intrinsic
binOp [PrimType]
anyIntType
    intrinsicBinOp BinOp
LogAnd = [PrimType] -> Maybe Intrinsic
binOp [PrimType
Bool]
    intrinsicBinOp BinOp
LogOr = [PrimType] -> Maybe Intrinsic
binOp [PrimType
Bool]
    intrinsicBinOp BinOp
Equal = Intrinsic -> Maybe Intrinsic
forall a. a -> Maybe a
Just Intrinsic
IntrinsicEquality
    intrinsicBinOp BinOp
NotEqual = Intrinsic -> Maybe Intrinsic
forall a. a -> Maybe a
Just Intrinsic
IntrinsicEquality
    intrinsicBinOp BinOp
Less = Maybe Intrinsic
ordering
    intrinsicBinOp BinOp
Leq = Maybe Intrinsic
ordering
    intrinsicBinOp BinOp
Greater = Maybe Intrinsic
ordering
    intrinsicBinOp BinOp
Geq = Maybe Intrinsic
ordering
    intrinsicBinOp BinOp
_ = Maybe Intrinsic
forall a. Maybe a
Nothing

-- | The largest tag used by an intrinsic - this can be used to
-- determine whether a 'VName' refers to an intrinsic or a user-defined name.
maxIntrinsicTag :: Int
maxIntrinsicTag :: Int
maxIntrinsicTag = [Int] -> Int
forall a (f :: * -> *). (Num a, Ord a, Foldable f) => f a -> a
maxinum ([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] -> [Int]) -> [VName] -> [Int]
forall a b. (a -> b) -> a -> b
$ Map VName Intrinsic -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName Intrinsic
intrinsics

-- | Create a name with no qualifiers from a name.
qualName :: v -> QualName v
qualName :: v -> QualName v
qualName = [v] -> v -> QualName v
forall vn. [vn] -> vn -> QualName vn
QualName []

-- | Add another qualifier (at the head) to a qualified name.
qualify :: v -> QualName v -> QualName v
qualify :: v -> QualName v -> QualName v
qualify v
k (QualName [v]
ks v
v) = [v] -> v -> QualName v
forall vn. [vn] -> vn -> QualName vn
QualName (v
k v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
ks) v
v

-- | Create a type name name with no qualifiers from a 'VName'.
typeName :: VName -> TypeName
typeName :: VName -> TypeName
typeName = QualName VName -> TypeName
typeNameFromQualName (QualName VName -> TypeName)
-> (VName -> QualName VName) -> VName -> TypeName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName

-- | The modules imported by a Futhark program.
progImports :: ProgBase f vn -> [(String, SrcLoc)]
progImports :: ProgBase f vn -> [(String, SrcLoc)]
progImports = (DecBase f vn -> [(String, SrcLoc)])
-> [DecBase f vn] -> [(String, SrcLoc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, SrcLoc)]
decImports ([DecBase f vn] -> [(String, SrcLoc)])
-> (ProgBase f vn -> [DecBase f vn])
-> ProgBase f vn
-> [(String, SrcLoc)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgBase f vn -> [DecBase f vn]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs

-- | The modules imported by a single declaration.
decImports :: DecBase f vn -> [(String, SrcLoc)]
decImports :: DecBase f vn -> [(String, SrcLoc)]
decImports (OpenDec ModExpBase f vn
x SrcLoc
_) = ModExpBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, SrcLoc)]
modExpImports ModExpBase f vn
x
decImports (ModDec ModBindBase f vn
md) = ModExpBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, SrcLoc)]
modExpImports (ModExpBase f vn -> [(String, SrcLoc)])
-> ModExpBase f vn -> [(String, SrcLoc)]
forall a b. (a -> b) -> a -> b
$ ModBindBase f vn -> ModExpBase f vn
forall (f :: * -> *) vn. ModBindBase f vn -> ModExpBase f vn
modExp ModBindBase f vn
md
decImports SigDec {} = []
decImports TypeDec {} = []
decImports ValDec {} = []
decImports (LocalDec DecBase f vn
d SrcLoc
_) = DecBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, SrcLoc)]
decImports DecBase f vn
d
decImports (ImportDec String
x f String
_ SrcLoc
loc) = [(String
x, SrcLoc
loc)]

modExpImports :: ModExpBase f vn -> [(String, SrcLoc)]
modExpImports :: ModExpBase f vn -> [(String, SrcLoc)]
modExpImports ModVar {} = []
modExpImports (ModParens ModExpBase f vn
p SrcLoc
_) = ModExpBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, SrcLoc)]
modExpImports ModExpBase f vn
p
modExpImports (ModImport String
f f String
_ SrcLoc
loc) = [(String
f, SrcLoc
loc)]
modExpImports (ModDecs [DecBase f vn]
ds SrcLoc
_) = (DecBase f vn -> [(String, SrcLoc)])
-> [DecBase f vn] -> [(String, SrcLoc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, SrcLoc)]
decImports [DecBase f vn]
ds
modExpImports (ModApply ModExpBase f vn
_ ModExpBase f vn
me f (Map VName VName)
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, SrcLoc)]
modExpImports ModExpBase f vn
me
modExpImports (ModAscript ModExpBase f vn
me SigExpBase f vn
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [(String, SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, SrcLoc)]
modExpImports ModExpBase f vn
me
modExpImports ModLambda {} = []

-- | The set of module types used in any exported (non-local)
-- declaration.
progModuleTypes :: Ord vn => ProgBase f vn -> S.Set vn
progModuleTypes :: ProgBase f vn -> Set vn
progModuleTypes = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn)
-> (ProgBase f vn -> [Set vn]) -> ProgBase f vn -> Set vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DecBase f vn -> Set vn) -> [DecBase f vn] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map DecBase f vn -> Set vn
forall (f :: * -> *). DecBase f vn -> Set vn
onDec ([DecBase f vn] -> [Set vn])
-> (ProgBase f vn -> [DecBase f vn]) -> ProgBase f vn -> [Set vn]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgBase f vn -> [DecBase f vn]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs
  where
    onDec :: DecBase f vn -> Set vn
onDec (OpenDec ModExpBase f vn
x SrcLoc
_) = ModExpBase f vn -> Set vn
onModExp ModExpBase f vn
x
    onDec (ModDec ModBindBase f vn
md) =
      Set vn
-> ((SigExpBase f vn, f (Map VName VName)) -> Set vn)
-> Maybe (SigExpBase f vn, f (Map VName VName))
-> Set vn
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set vn
forall a. Monoid a => a
mempty (SigExpBase f vn -> Set vn
forall a (f :: * -> *). Ord a => SigExpBase f a -> Set a
onSigExp (SigExpBase f vn -> Set vn)
-> ((SigExpBase f vn, f (Map VName VName)) -> SigExpBase f vn)
-> (SigExpBase f vn, f (Map VName VName))
-> Set vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SigExpBase f vn, f (Map VName VName)) -> SigExpBase f vn
forall a b. (a, b) -> a
fst) (ModBindBase f vn -> Maybe (SigExpBase f vn, f (Map VName VName))
forall (f :: * -> *) vn.
ModBindBase f vn -> Maybe (SigExpBase f vn, f (Map VName VName))
modSignature ModBindBase f vn
md) Set vn -> Set vn -> Set vn
forall a. Semigroup a => a -> a -> a
<> ModExpBase f vn -> Set vn
onModExp (ModBindBase f vn -> ModExpBase f vn
forall (f :: * -> *) vn. ModBindBase f vn -> ModExpBase f vn
modExp ModBindBase f vn
md)
    onDec SigDec {} = Set vn
forall a. Monoid a => a
mempty
    onDec TypeDec {} = Set vn
forall a. Monoid a => a
mempty
    onDec ValDec {} = Set vn
forall a. Monoid a => a
mempty
    onDec LocalDec {} = Set vn
forall a. Monoid a => a
mempty
    onDec ImportDec {} = Set vn
forall a. Monoid a => a
mempty

    onModExp :: ModExpBase f vn -> Set vn
onModExp ModVar {} = Set vn
forall a. Monoid a => a
mempty
    onModExp (ModParens ModExpBase f vn
p SrcLoc
_) = ModExpBase f vn -> Set vn
onModExp ModExpBase f vn
p
    onModExp ModImport {} = Set vn
forall a. Monoid a => a
mempty
    onModExp (ModDecs [DecBase f vn]
ds SrcLoc
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ (DecBase f vn -> Set vn) -> [DecBase f vn] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map DecBase f vn -> Set vn
onDec [DecBase f vn]
ds
    onModExp (ModApply ModExpBase f vn
me1 ModExpBase f vn
me2 f (Map VName VName)
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> Set vn
onModExp ModExpBase f vn
me1 Set vn -> Set vn -> Set vn
forall a. Semigroup a => a -> a -> a
<> ModExpBase f vn -> Set vn
onModExp ModExpBase f vn
me2
    onModExp (ModAscript ModExpBase f vn
me SigExpBase f vn
se f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> Set vn
onModExp ModExpBase f vn
me Set vn -> Set vn -> Set vn
forall a. Semigroup a => a -> a -> a
<> SigExpBase f vn -> Set vn
forall a (f :: * -> *). Ord a => SigExpBase f a -> Set a
onSigExp SigExpBase f vn
se
    onModExp (ModLambda ModParamBase f vn
p Maybe (SigExpBase f vn, f (Map VName VName))
r ModExpBase f vn
me SrcLoc
_) =
      ModParamBase f vn -> Set vn
forall (f :: * -> *). ModParamBase f vn -> Set vn
onModParam ModParamBase f vn
p Set vn -> Set vn -> Set vn
forall a. Semigroup a => a -> a -> a
<> Set vn
-> ((SigExpBase f vn, f (Map VName VName)) -> Set vn)
-> Maybe (SigExpBase f vn, f (Map VName VName))
-> Set vn
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set vn
forall a. Monoid a => a
mempty (SigExpBase f vn -> Set vn
forall a (f :: * -> *). Ord a => SigExpBase f a -> Set a
onSigExp (SigExpBase f vn -> Set vn)
-> ((SigExpBase f vn, f (Map VName VName)) -> SigExpBase f vn)
-> (SigExpBase f vn, f (Map VName VName))
-> Set vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SigExpBase f vn, f (Map VName VName)) -> SigExpBase f vn
forall a b. (a, b) -> a
fst) Maybe (SigExpBase f vn, f (Map VName VName))
r Set vn -> Set vn -> Set vn
forall a. Semigroup a => a -> a -> a
<> ModExpBase f vn -> Set vn
onModExp ModExpBase f vn
me

    onModParam :: ModParamBase f vn -> Set vn
onModParam = SigExpBase f vn -> Set vn
forall a (f :: * -> *). Ord a => SigExpBase f a -> Set a
onSigExp (SigExpBase f vn -> Set vn)
-> (ModParamBase f vn -> SigExpBase f vn)
-> ModParamBase f vn
-> Set vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModParamBase f vn -> SigExpBase f vn
forall (f :: * -> *) vn. ModParamBase f vn -> SigExpBase f vn
modParamType

    onSigExp :: SigExpBase f a -> Set a
onSigExp (SigVar QualName a
v f (Map VName VName)
_ SrcLoc
_) = a -> Set a
forall a. a -> Set a
S.singleton (a -> Set a) -> a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
v
    onSigExp (SigParens SigExpBase f a
e SrcLoc
_) = SigExpBase f a -> Set a
onSigExp SigExpBase f a
e
    onSigExp SigSpecs {} = Set a
forall a. Monoid a => a
mempty
    onSigExp (SigWith SigExpBase f a
e TypeRefBase f a
_ SrcLoc
_) = SigExpBase f a -> Set a
onSigExp SigExpBase f a
e
    onSigExp (SigArrow Maybe a
_ SigExpBase f a
e1 SigExpBase f a
e2 SrcLoc
_) = SigExpBase f a -> Set a
onSigExp SigExpBase f a
e1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> SigExpBase f a -> Set a
onSigExp SigExpBase f a
e2

-- | Extract a leading @((name, namespace, file), remainder)@ from a
-- documentation comment string.  These are formatted as
-- \`name\`\@namespace[\@file].  Let us hope that this pattern does not occur
-- anywhere else.
identifierReference :: String -> Maybe ((String, String, Maybe FilePath), String)
identifierReference :: String -> Maybe ((String, String, Maybe String), String)
identifierReference (Char
'`' : String
s)
  | (String
identifier, Char
'`' : Char
'@' : String
s') <- (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'`') String
s,
    (String
namespace, String
s'') <- (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isAlpha String
s',
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
namespace =
    case String
s'' of
      Char
'@' : Char
'"' : String
s'''
        | (String
file, Char
'"' : String
s'''') <- (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'"') String
s''' ->
          ((String, String, Maybe String), String)
-> Maybe ((String, String, Maybe String), String)
forall a. a -> Maybe a
Just ((String
identifier, String
namespace, String -> Maybe String
forall a. a -> Maybe a
Just String
file), String
s'''')
      String
_ -> ((String, String, Maybe String), String)
-> Maybe ((String, String, Maybe String), String)
forall a. a -> Maybe a
Just ((String
identifier, String
namespace, Maybe String
forall a. Maybe a
Nothing), String
s'')
identifierReference String
_ = Maybe ((String, String, Maybe String), String)
forall a. Maybe a
Nothing

-- | Given an operator name, return the operator that determines its
-- syntactical properties.
leadingOperator :: Name -> BinOp
leadingOperator :: Name -> BinOp
leadingOperator Name
s =
  BinOp
-> ((String, BinOp) -> BinOp) -> Maybe (String, BinOp) -> BinOp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe BinOp
Backtick (String, BinOp) -> BinOp
forall a b. (a, b) -> b
snd (Maybe (String, BinOp) -> BinOp) -> Maybe (String, BinOp) -> BinOp
forall a b. (a -> b) -> a -> b
$
    ((String, BinOp) -> Bool)
-> [(String, BinOp)] -> Maybe (String, BinOp)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` String
s') (String -> Bool)
-> ((String, BinOp) -> String) -> (String, BinOp) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, BinOp) -> String
forall a b. (a, b) -> a
fst) ([(String, BinOp)] -> Maybe (String, BinOp))
-> [(String, BinOp)] -> Maybe (String, BinOp)
forall a b. (a -> b) -> a -> b
$
      ((String, BinOp) -> Down Int)
-> [(String, BinOp)] -> [(String, BinOp)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int -> Down Int
forall a. a -> Down a
Down (Int -> Down Int)
-> ((String, BinOp) -> Int) -> (String, BinOp) -> Down Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (String -> Int)
-> ((String, BinOp) -> String) -> (String, BinOp) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, BinOp) -> String
forall a b. (a, b) -> a
fst) ([(String, BinOp)] -> [(String, BinOp)])
-> [(String, BinOp)] -> [(String, BinOp)]
forall a b. (a -> b) -> a -> b
$
        [String] -> [BinOp] -> [(String, BinOp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((BinOp -> String) -> [BinOp] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map BinOp -> String
forall a. Pretty a => a -> String
pretty [BinOp]
operators) [BinOp]
operators
  where
    s' :: String
s' = Name -> String
nameToString Name
s
    operators :: [BinOp]
    operators :: [BinOp]
operators = [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound :: BinOp]

-- | A type with no aliasing information but shape annotations.
type UncheckedType = TypeBase (ShapeDecl Name) ()

-- | An expression with no type annotations.
type UncheckedTypeExp = TypeExp Name

-- | A type declaration with no expanded type.
type UncheckedTypeDecl = TypeDeclBase NoInfo Name

-- | An identifier with no type annotations.
type UncheckedIdent = IdentBase NoInfo Name

-- | An index with no type annotations.
type UncheckedDimIndex = DimIndexBase NoInfo Name

-- | An expression with no type annotations.
type UncheckedExp = ExpBase NoInfo Name

-- | A module expression with no type annotations.
type UncheckedModExp = ModExpBase NoInfo Name

-- | A module type expression with no type annotations.
type UncheckedSigExp = SigExpBase NoInfo Name

-- | A type parameter with no type annotations.
type UncheckedTypeParam = TypeParamBase Name

-- | A pattern with no type annotations.
type UncheckedPattern = PatternBase NoInfo Name

-- | A function declaration with no type annotations.
type UncheckedValBind = ValBindBase NoInfo Name

-- | A declaration with no type annotations.
type UncheckedDec = DecBase NoInfo Name

-- | A spec with no type annotations.
type UncheckedSpec = SpecBase NoInfo Name

-- | A Futhark program with no type annotations.
type UncheckedProg = ProgBase NoInfo Name

-- | A case (of a match expression) with no type annotations.
type UncheckedCase = CaseBase NoInfo Name