-- | 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,
    isBuiltin,
    isBuiltinLoc,
    maxIntrinsicTag,
    namesToPrimTypes,
    qualName,
    qualify,
    primValueType,
    leadingOperator,
    progImports,
    decImports,
    progModuleTypes,
    identifierReference,
    prettyStacktrace,
    progHoles,
    defaultEntryPoint,
    paramName,

    -- * Queries on expressions
    typeOf,
    valBindTypeScheme,
    valBindBound,
    funType,

    -- * Queries on patterns and params
    patIdents,
    patNames,
    patternMap,
    patternType,
    patternStructType,
    patternParam,
    patternOrderZero,

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

    -- * Operations on types
    peelArray,
    stripArray,
    arrayOf,
    toStructural,
    toStruct,
    fromStruct,
    setAliases,
    addAliases,
    setUniqueness,
    noSizes,
    traverseDims,
    DimPos (..),
    tupleRecord,
    isTupleRecord,
    areTupleFields,
    tupleFields,
    tupleFieldNames,
    sortFields,
    sortConstrs,
    isTypeParam,
    isSizeParam,
    combineTypeShapes,
    matchDims,
    -- | 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,
    UncheckedDimIndex,
    UncheckedSlice,
    UncheckedExp,
    UncheckedModExp,
    UncheckedSigExp,
    UncheckedTypeParam,
    UncheckedPat,
    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, sortOn)
import Data.Loc (Loc (..), posFile)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Data.Set qualified as S
import Futhark.Util (maxinum)
import Futhark.Util.Pretty
import Language.Futhark.Primitive qualified as Primitive
import Language.Futhark.Syntax
import Language.Futhark.Traversals
import Language.Futhark.Tuple
import System.FilePath (takeDirectory)

-- | The name of the default program entry point (@main@).
defaultEntryPoint :: Name
defaultEntryPoint :: Name
defaultEntryPoint = FilePath -> Name
nameFromString FilePath
"main"

-- | 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 :: forall dim as. TypeBase dim as -> Int
arrayRank = forall dim. Shape dim -> Int
shapeRank forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. TypeBase dim as -> Shape dim
arrayShape

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

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

-- | 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
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
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
Ord, Int -> DimPos -> ShowS
[DimPos] -> ShowS
DimPos -> FilePath
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [DimPos] -> ShowS
$cshowList :: [DimPos] -> ShowS
show :: DimPos -> FilePath
$cshow :: DimPos -> FilePath
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 :: forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName -> DimPos -> fdim -> f tdim
f = forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go forall a. Monoid a => a
mempty DimPos
PosImmediate
  where
    go ::
      forall als'.
      S.Set VName ->
      DimPos ->
      TypeBase fdim als' ->
      f (TypeBase tdim als')
    go :: forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b t :: TypeBase fdim als'
t@Array {} =
      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) 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)) =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (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 QualName VName
tn [TypeArg fdim]
targs)) =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar als'
as Uniqueness
u QualName VName
tn forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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)) =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (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)) =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    go Set VName
bound DimPos
_ (Scalar (Arrow als'
als PName
p TypeBase fdim ()
t1 (RetType [VName]
dims TypeBase fdim als'
t2))) =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow als'
als PName
p forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosParam TypeBase fdim ()
t1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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' =
          forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims
            forall a. Semigroup a => a -> a -> a
<> case PName
p of
              Named VName
p' -> 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) =
      forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim 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 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    onTypeArg Set VName
bound DimPos
b (TypeArgType TypeBase fdim ()
t SrcLoc
loc) =
      forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b TypeBase fdim ()
t forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

-- | Return the uniqueness of a type.
uniqueness :: TypeBase shape as -> Uniqueness
uniqueness :: forall shape as. TypeBase shape as -> Uniqueness
uniqueness (Array as
_ Uniqueness
u Shape shape
_ ScalarTypeBase shape ()
_) = Uniqueness
u
uniqueness (Scalar (TypeVar as
_ Uniqueness
u QualName VName
_ [TypeArg shape]
_)) = Uniqueness
u
uniqueness (Scalar (Sum Map Name [TypeBase shape as]
ts))
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any forall shape as. TypeBase shape as -> Bool
unique) Map Name [TypeBase shape as]
ts = Uniqueness
Unique
uniqueness (Scalar (Record Map Name (TypeBase shape as)
fs))
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any forall shape as. TypeBase shape as -> Bool
unique Map Name (TypeBase shape as)
fs = Uniqueness
Unique
uniqueness TypeBase shape as
_ = Uniqueness
Nonunique

-- | @unique t@ is 'True' if the type of the argument is unique.
unique :: TypeBase shape as -> Bool
unique :: forall shape as. TypeBase shape as -> Bool
unique = (forall a. Eq a => a -> a -> Bool
== Uniqueness
Unique) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall as shape. Monoid as => TypeBase shape as -> as
aliases = forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap (forall a b. a -> b -> a
const forall a. Monoid a => a
mempty) 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 :: forall shape as. TypeBase shape as -> Diet
diet (Scalar (Record Map Name (TypeBase shape as)
ets)) = Map Name Diet -> Diet
RecordDiet forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap 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 ()
t1 (RetType [VName]
_ TypeBase shape as
t2))) = Diet -> Diet -> Diet
FuncDiet (forall shape as. TypeBase shape as -> Diet
diet TypeBase shape ()
t1) (forall shape as. TypeBase shape as -> Diet
diet TypeBase shape as
t2)
diet (Array as
_ Uniqueness
Unique Shape shape
_ ScalarTypeBase shape ()
_) = Diet
Consume
diet (Array as
_ Uniqueness
Nonunique Shape shape
_ ScalarTypeBase shape ()
_) = Diet
Observe
diet (Scalar (TypeVar as
_ Uniqueness
Unique QualName VName
_ [TypeArg shape]
_)) = Diet
Consume
diet (Scalar (TypeVar as
_ Uniqueness
Nonunique QualName VName
_ [TypeArg shape]
_)) = Diet
Observe
diet (Scalar (Sum Map Name [TypeBase shape as]
cs)) = Map Name [Diet] -> Diet
SumDiet forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall shape as. TypeBase shape as -> Diet
diet) Map Name [TypeBase shape as]
cs

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

-- | Remove aliasing information from a type.
toStruct ::
  TypeBase dim as ->
  TypeBase dim ()
toStruct :: forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase dim as
t = TypeBase dim as
t 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 :: forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase dim as
t = TypeBase dim as
t forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` 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 :: forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray Int
n (Array as
als Uniqueness
u Shape dim
shape ScalarTypeBase dim ()
t)
  | forall dim. Shape dim -> Int
shapeRank Shape dim
shape forall a. Eq a => a -> a -> Bool
== Int
n =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
t forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` forall a b. a -> b -> a
const as
als
  | Bool
otherwise =
      forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array as
als Uniqueness
u forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall dim. Int -> Shape dim -> Maybe (Shape dim)
stripDims Int
n Shape dim
shape forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase dim ()
t
peelArray Int
_ TypeBase dim as
_ = forall a. Maybe a
Nothing

-- | @arrayOf u s t@ 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 =>
  Uniqueness ->
  Shape dim ->
  TypeBase dim as ->
  TypeBase dim as
arrayOf :: forall as dim.
Monoid as =>
Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOf Uniqueness
u Shape dim
s TypeBase dim as
t = forall as dim.
Monoid as =>
as -> Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOfWithAliases forall a. Monoid a => a
mempty Uniqueness
u Shape dim
s (TypeBase dim as
t forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique)

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

-- | @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 :: forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
n (Array as
als Uniqueness
u Shape dim
shape ScalarTypeBase dim ()
et)
  | Just Shape dim
shape' <- forall dim. Int -> Shape dim -> Maybe (Shape dim)
stripDims Int
n Shape dim
shape =
      forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array as
als Uniqueness
u Shape dim
shape' ScalarTypeBase dim ()
et
  | Bool
otherwise =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
et forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u 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] -> ScalarTypeBase dim as
tupleRecord :: forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord = forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord (Scalar (Record Map Name (TypeBase dim as)
fs)) = forall a. Map Name a -> Maybe [a]
areTupleFields Map Name (TypeBase dim as)
fs
isTupleRecord TypeBase dim as
_ = forall a. Maybe a
Nothing

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

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

-- | Is this a 'TypeParamDim'?
isSizeParam :: TypeParamBase vn -> Bool
isSizeParam :: forall vn. TypeParamBase vn -> Bool
isSizeParam = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) =>
  TypeBase Size as ->
  TypeBase Size as ->
  TypeBase Size as
combineTypeShapes :: forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes (Scalar (Record Map Name (TypeBase Size as)
ts1)) (Scalar (Record Map Name (TypeBase Size as)
ts2))
  | forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Size as)
ts1 forall a. Eq a => a -> a -> Bool
== forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Size as)
ts2 =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$
        forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$
          forall a b k. (a -> b) -> Map k a -> Map k b
M.map
            (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes)
            (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name (TypeBase Size as)
ts1 Map Name (TypeBase Size as)
ts2)
combineTypeShapes (Scalar (Sum Map Name [TypeBase Size as]
cs1)) (Scalar (Sum Map Name [TypeBase Size as]
cs2))
  | forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Size as]
cs1 forall a. Eq a => a -> a -> Bool
== forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Size as]
cs2 =
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$
        forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$
          forall a b k. (a -> b) -> Map k a -> Map k b
M.map
            (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes)
            (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name [TypeBase Size as]
cs1 Map Name [TypeBase Size as]
cs2)
combineTypeShapes (Scalar (Arrow as
als1 PName
p1 StructType
a1 (RetType [VName]
dims1 TypeBase Size as
b1))) (Scalar (Arrow as
als2 PName
_p2 StructType
a2 (RetType [VName]
_ TypeBase Size as
b2))) =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow (as
als1 forall a. Semigroup a => a -> a -> a
<> as
als2) PName
p1 (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes StructType
a1 StructType
a2) (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims1 (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes TypeBase Size as
b1 TypeBase Size as
b2))
combineTypeShapes (Scalar (TypeVar as
als1 Uniqueness
u1 QualName VName
v [TypeArg Size]
targs1)) (Scalar (TypeVar as
als2 Uniqueness
_ QualName VName
_ [TypeArg Size]
targs2)) =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar (as
als1 forall a. Semigroup a => a -> a -> a
<> as
als2) Uniqueness
u1 QualName VName
v forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeArg Size -> TypeArg Size -> TypeArg Size
f [TypeArg Size]
targs1 [TypeArg Size]
targs2
  where
    f :: TypeArg Size -> TypeArg Size -> TypeArg Size
f (TypeArgType StructType
t1 SrcLoc
loc) (TypeArgType StructType
t2 SrcLoc
_) =
      forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes StructType
t1 StructType
t2) SrcLoc
loc
    f TypeArg Size
targ TypeArg Size
_ = TypeArg Size
targ
combineTypeShapes (Array as
als1 Uniqueness
u1 Shape Size
shape1 ScalarTypeBase Size ()
et1) (Array as
als2 Uniqueness
_u2 Shape Size
_shape2 ScalarTypeBase Size ()
et2) =
  forall as dim.
Monoid as =>
as -> Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOfWithAliases
    (as
als1 forall a. Semigroup a => a -> a -> a
<> as
als2)
    Uniqueness
u1
    Shape Size
shape1
    (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase Size ()
et1) (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase Size ()
et2) forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` forall a. Monoid a => a
mempty)
combineTypeShapes TypeBase Size as
_ TypeBase Size as
new_tp = TypeBase Size as
new_tp

-- | The name, if any.
paramName :: PName -> Maybe VName
paramName :: PName -> Maybe VName
paramName (Named VName
v) = forall a. a -> Maybe a
Just VName
v
paramName PName
Unnamed = forall a. Maybe a
Nothing

-- | Match the dimensions of otherwise assumed-equal types.  The
-- combining function is also passed the names bound within the type
-- (from named parameters or return types).
matchDims ::
  forall as m d1 d2.
  (Monoid as, Monad m) =>
  ([VName] -> d1 -> d2 -> m d1) ->
  TypeBase d1 as ->
  TypeBase d2 as ->
  m (TypeBase d1 as)
matchDims :: forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims [VName] -> d1 -> d2 -> m d1
onDims = forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' forall a. Monoid a => a
mempty
  where
    matchDims' ::
      forall as'. Monoid as' => [VName] -> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
    matchDims' :: forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound TypeBase d1 as'
t1 TypeBase d2 as'
t2 =
      case (TypeBase d1 as'
t1, TypeBase d2 as'
t2) of
        (Array as'
als1 Uniqueness
u1 Shape d1
shape1 ScalarTypeBase d1 ()
et1, Array as'
als2 Uniqueness
u2 Shape d2
shape2 ScalarTypeBase d2 ()
et2) ->
          forall a b c. (a -> b -> c) -> b -> a -> c
flip forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases (as'
als1 forall a. Semigroup a => a -> a -> a
<> as'
als2)
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( forall as dim.
Monoid as =>
Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOf (forall a. Ord a => a -> a -> a
min Uniqueness
u1 Uniqueness
u2)
                    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName] -> Shape d1 -> Shape d2 -> m (Shape d1)
onShapes [VName]
bound Shape d1
shape1 Shape d2
shape2
                    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d1 ()
et1) (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d2 ()
et2)
                )
        (Scalar (Record Map Name (TypeBase d1 as')
f1), Scalar (Record Map Name (TypeBase d2 as')
f2)) ->
          forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound)) (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)) ->
          forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
              (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound)))
              (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith 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 ()
a1 (RetType [VName]
dims1 TypeBase d1 as'
b1)),
          Scalar (Arrow as'
als2 PName
p2 TypeBase d2 ()
a2 (RetType [VName]
dims2 TypeBase d2 as'
b2))
          ) ->
            let bound' :: [VName]
bound' = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe PName -> Maybe VName
paramName [PName
p1, PName
p2] forall a. Semigroup a => a -> a -> a
<> [VName]
dims1 forall a. Semigroup a => a -> a -> a
<> [VName]
dims2 forall a. Semigroup a => a -> a -> a
<> [VName]
bound
             in forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar
                  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow (as'
als1 forall a. Semigroup a => a -> a -> a
<> as'
als2) PName
p1
                          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound' TypeBase d1 ()
a1 TypeBase d2 ()
a2
                          forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound' TypeBase d1 as'
b1 TypeBase d2 as'
b2)
                      )
        ( Scalar (TypeVar as'
als1 Uniqueness
u QualName VName
v [TypeArg d1]
targs1),
          Scalar (TypeVar as'
als2 Uniqueness
_ QualName VName
_ [TypeArg d2]
targs2)
          ) ->
            forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar (as'
als1 forall a. Semigroup a => a -> a -> a
<> as'
als2) Uniqueness
u QualName VName
v
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ([VName] -> TypeArg d1 -> TypeArg d2 -> m (TypeArg d1)
matchTypeArg [VName]
bound) [TypeArg d1]
targs1 [TypeArg d2]
targs2
        (TypeBase d1 as', TypeBase d2 as')
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase d1 as'
t1

    matchTypeArg :: [VName] -> TypeArg d1 -> TypeArg d2 -> m (TypeArg d1)
matchTypeArg [VName]
_ ta :: TypeArg d1
ta@TypeArgType {} TypeArg d2
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeArg d1
ta
    matchTypeArg [VName]
bound (TypeArgDim d1
x SrcLoc
loc) (TypeArgDim d2
y SrcLoc
_) =
      forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName] -> d1 -> d2 -> m d1
onDims [VName]
bound d1
x d2
y forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    matchTypeArg [VName]
_ TypeArg d1
a TypeArg d2
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeArg d1
a

    onShapes :: [VName] -> Shape d1 -> Shape d2 -> m (Shape d1)
onShapes [VName]
bound Shape d1
shape1 Shape d2
shape2 =
      forall dim. [dim] -> Shape dim
Shape forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ([VName] -> d1 -> d2 -> m d1
onDims [VName]
bound) (forall dim. Shape dim -> [dim]
shapeDims Shape d1
shape1) (forall dim. Shape dim -> [dim]
shapeDims Shape 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 :: forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
setUniqueness (Array as
als Uniqueness
_ Shape dim
shape ScalarTypeBase dim ()
et) Uniqueness
u =
  forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array as
als Uniqueness
u Shape dim
shape ScalarTypeBase dim ()
et
setUniqueness (Scalar (TypeVar as
als Uniqueness
_ QualName VName
t [TypeArg dim]
targs)) Uniqueness
u =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar as
als Uniqueness
u QualName VName
t [TypeArg dim]
targs
setUniqueness (Scalar (Record Map Name (TypeBase dim as)
ets)) Uniqueness
u =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (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 =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map (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 :: forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases TypeBase dim asf
t = forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
addAliases TypeBase dim asf
t forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
addAliases = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 Float16Value {} = FloatType
Float16
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 forall a b. (a -> b) -> a -> b
$ IntValue -> IntType
intValueType IntValue
v
primValueType (UnsignedValue IntValue
v) = IntType -> PrimType
Unsigned forall a b. (a -> b) -> a -> b
$ IntValue -> IntType
intValueType IntValue
v
primValueType (FloatValue FloatValue
v) = FloatType -> PrimType
FloatType forall a b. (a -> b) -> a -> b
$ FloatValue -> FloatType
floatValueType FloatValue
v
primValueType BoolValue {} = PrimType
Bool

-- | 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 -> PatType
typeOf :: ExpBase Info VName -> PatType
typeOf (Literal PrimValue
val SrcLoc
_) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
typeOf (IntLit Integer
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (FloatLit Double
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Parens ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (QualParens (QualName VName, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (TupLit [ExpBase Info VName]
es SrcLoc
_) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map ExpBase Info VName -> PatType
typeOf [ExpBase Info VName]
es
typeOf (RecordLit [FieldBase Info VName]
fs SrcLoc
_) =
  -- Reverse, because M.unions is biased to the left.
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map FieldBase Info VName -> Map Name PatType
record [FieldBase Info VName]
fs
  where
    record :: FieldBase Info VName -> Map Name PatType
record (RecordFieldExplicit Name
name ExpBase Info VName
e SrcLoc
_) = forall k a. k -> a -> Map k a
M.singleton Name
name forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
    record (RecordFieldImplicit VName
name (Info PatType
t) SrcLoc
_) =
      forall k a. k -> a -> Map k a
M.singleton (VName -> Name
baseName VName
name) forall a b. (a -> b) -> a -> b
$
        PatType
t
          forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` forall a. Ord a => a -> Set a -> Set a
S.insert (VName -> Alias
AliasBound VName
name)
typeOf (ArrayLit [ExpBase Info VName]
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (StringLit [Word8]
vs SrcLoc
_) =
  forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array
    forall a. Monoid a => a
mempty
    Uniqueness
Unique
    (forall dim. [dim] -> Shape dim
Shape [Int -> Size
ConstSize forall a b. (a -> b) -> a -> b
$ forall i a. Num i => [a] -> i
genericLength [Word8]
vs])
    (forall dim as. PrimType -> ScalarTypeBase dim as
Prim (IntType -> PrimType
Unsigned IntType
Int8))
typeOf (Project Name
_ ExpBase Info VName
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Var QualName VName
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Hole (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Ascript ExpBase Info VName
e TypeExp VName
_ SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (Negate ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (Not ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (Update ExpBase Info VName
e SliceBase Info VName
_ ExpBase Info VName
_ SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` forall a. Monoid a => a
mempty
typeOf (RecordUpdate ExpBase Info VName
_ [Name]
_ ExpBase Info VName
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Assert ExpBase Info VName
_ ExpBase Info VName
e Info Text
_ SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (Lambda [PatBase Info VName]
params ExpBase Info VName
_ Maybe (TypeExp VName)
_ (Info (Aliasing
als, StructRetType
t)) SrcLoc
_) =
  let RetType [] StructType
t' = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall {dim}.
(PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> (PName, StructType)
patternParam) StructRetType
t [PatBase Info VName]
params
   in StructType
t' forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliasing
als
  where
    arrow :: (PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow (Named VName
v, TypeBase dim ()
x) (RetType [VName]
dims TypeBase dim ()
y) =
      forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () (VName -> PName
Named VName
v) TypeBase dim ()
x forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (VName
v forall a. a -> [a] -> [a]
: [VName]
dims) TypeBase dim ()
y
    arrow (PName
pn, TypeBase dim ()
tx) RetTypeBase dim ()
y =
      forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
pn TypeBase dim ()
tx RetTypeBase dim ()
y
typeOf (OpSection QualName VName
_ (Info PatType
t) SrcLoc
_) =
  PatType
t
typeOf (OpSectionLeft QualName VName
_ Info PatType
_ ExpBase Info VName
_ (Info (PName, StructType, Maybe VName)
_, Info (PName
pn, StructType
pt2)) (Info PatRetType
ret, Info [VName]
_) SrcLoc
_) =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
pn StructType
pt2 PatRetType
ret
typeOf (OpSectionRight QualName VName
_ Info PatType
_ ExpBase Info VName
_ (Info (PName
pn, StructType
pt1), Info (PName, StructType, Maybe VName)
_) (Info PatRetType
ret) SrcLoc
_) =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
pn StructType
pt1 PatRetType
ret
typeOf (ProjectSection [Name]
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (IndexSection SliceBase Info VName
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Constr Name
_ [ExpBase Info VName]
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Attr AttrInfo VName
_ ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
typeOf (AppExp AppExpBase Info VName
_ (Info AppRes
res)) = AppRes -> PatType
appResType AppRes
res

-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes
-- @ts@ as parameters and returns @ret@.
foldFunType ::
  Monoid as =>
  [TypeBase dim pas] ->
  RetTypeBase dim as ->
  TypeBase dim as
foldFunType :: forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [TypeBase dim pas]
ps RetTypeBase dim as
ret =
  let RetType [VName]
_ TypeBase dim as
t = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {as} {dim} {as}.
Monoid as =>
TypeBase dim as -> RetTypeBase dim as -> RetTypeBase dim as
arrow RetTypeBase dim as
ret [TypeBase dim pas]
ps
   in TypeBase dim as
t
  where
    arrow :: TypeBase dim as -> RetTypeBase dim as -> RetTypeBase dim as
arrow TypeBase dim as
t1 RetTypeBase dim as
t2 =
      forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase dim as
t1) RetTypeBase 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 ()], TypeBase dim ())
unfoldFunType :: forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType (Scalar (Arrow as
_ PName
_ TypeBase dim ()
t1 (RetType [VName]
_ TypeBase dim as
t2))) =
  let ([TypeBase dim ()]
ps, TypeBase dim ()
r) = forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType TypeBase dim as
t2
   in (TypeBase dim ()
t1 forall a. a -> [a] -> [a]
: [TypeBase dim ()]
ps, TypeBase dim ()
r)
unfoldFunType TypeBase dim as
t = ([], forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase dim as
t)

-- | The type scheme of a value binding, comprising the type
-- parameters and the actual type.
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBindBase Info VName
vb =
  ( forall (f :: * -> *) vn. ValBindBase f vn -> [TypeParamBase vn]
valBindTypeParams ValBindBase Info VName
vb,
    [PatBase Info VName] -> StructRetType -> StructType
funType (forall (f :: * -> *) vn. ValBindBase f vn -> [PatBase f vn]
valBindParams ValBindBase Info VName
vb) (forall a. Info a -> a
unInfo (forall (f :: * -> *) vn. ValBindBase f vn -> f StructRetType
valBindRetType ValBindBase Info VName
vb))
  )

-- | The names that are brought into scope by this value binding (not
-- including its own parameter names, but including any existential
-- sizes).
valBindBound :: ValBindBase Info VName -> [VName]
valBindBound :: ValBindBase Info VName -> [VName]
valBindBound ValBindBase Info VName
vb =
  forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBindBase Info VName
vb
    forall a. a -> [a] -> [a]
: case forall (f :: * -> *) vn. ValBindBase f vn -> [PatBase f vn]
valBindParams ValBindBase Info VName
vb of
      [] -> forall dim as. RetTypeBase dim as -> [VName]
retDims (forall a. Info a -> a
unInfo (forall (f :: * -> *) vn. ValBindBase f vn -> f StructRetType
valBindRetType ValBindBase Info VName
vb))
      [PatBase Info VName]
_ -> []

-- | The type of a function with the given parameters and return type.
funType :: [PatBase Info VName] -> StructRetType -> StructType
funType :: [PatBase Info VName] -> StructRetType -> StructType
funType [PatBase Info VName]
params StructRetType
ret =
  let RetType [VName]
_ StructType
t = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall {dim}.
(PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> (PName, StructType)
patternParam) StructRetType
ret [PatBase Info VName]
params
   in StructType
t
  where
    arrow :: (PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow (PName
xp, TypeBase dim ()
xt) RetTypeBase dim ()
yt = forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
xp TypeBase dim ()
xt RetTypeBase dim ()
yt

-- | The type names mentioned in a type.
typeVars :: Monoid as => TypeBase dim as -> S.Set VName
typeVars :: forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim as
t =
  case TypeBase dim as
t of
    Scalar Prim {} -> forall a. Monoid a => a
mempty
    Scalar (TypeVar as
_ Uniqueness
_ QualName VName
tn [TypeArg dim]
targs) ->
      forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a. a -> Set a
S.singleton (forall vn. QualName vn -> vn
qualLeaf QualName VName
tn) forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall {dim}. TypeArg dim -> Set VName
typeArgFree [TypeArg dim]
targs
    Scalar (Arrow as
_ PName
_ TypeBase dim ()
t1 (RetType [VName]
_ TypeBase dim as
t2)) -> forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim ()
t1 forall a. Semigroup a => a -> a -> a
<> forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim as
t2
    Scalar (Record Map Name (TypeBase dim as)
fields) -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap 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) -> forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars Map Name [TypeBase dim as]
cs
    Array as
_ Uniqueness
_ Shape dim
_ ScalarTypeBase dim ()
rt -> forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
rt
  where
    typeArgFree :: TypeArg dim -> Set VName
typeArgFree (TypeArgType TypeBase dim ()
ta SrcLoc
_) = forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim ()
ta
    typeArgFree TypeArgDim {} = 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 :: forall shape as. TypeBase shape as -> Bool
orderZero Array {} = Bool
True
orderZero (Scalar (Prim PrimType
_)) = Bool
True
orderZero (Scalar (Record Map Name (TypeBase dim as)
fs)) = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall shape as. TypeBase shape as -> Bool
orderZero forall a b. (a -> b) -> a -> b
$ 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)) = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall shape as. TypeBase shape as -> Bool
orderZero) Map Name [TypeBase dim as]
cs

-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern
-- have order 0.
patternOrderZero :: PatBase Info vn -> Bool
patternOrderZero :: forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info vn
pat = case PatBase Info vn
pat of
  TuplePat [PatBase Info vn]
ps SrcLoc
_ -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall vn. PatBase Info vn -> Bool
patternOrderZero [PatBase Info vn]
ps
  RecordPat [(Name, PatBase Info vn)]
fs SrcLoc
_ -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall vn. PatBase Info vn -> Bool
patternOrderZero forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, PatBase Info vn)]
fs
  PatParens PatBase Info vn
p SrcLoc
_ -> forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info vn
p
  Id vn
_ (Info PatType
t) SrcLoc
_ -> forall shape as. TypeBase shape as -> Bool
orderZero PatType
t
  Wildcard (Info PatType
t) SrcLoc
_ -> forall shape as. TypeBase shape as -> Bool
orderZero PatType
t
  PatAscription PatBase Info vn
p TypeExp vn
_ SrcLoc
_ -> forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info vn
p
  PatLit PatLit
_ (Info PatType
t) SrcLoc
_ -> forall shape as. TypeBase shape as -> Bool
orderZero PatType
t
  PatConstr Name
_ Info PatType
_ [PatBase Info vn]
ps SrcLoc
_ -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall vn. PatBase Info vn -> Bool
patternOrderZero [PatBase Info vn]
ps
  PatAttr AttrInfo vn
_ PatBase Info vn
p SrcLoc
_ -> forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info vn
p

-- | The set of identifiers bound in a pattern.
patIdents :: (Functor f, Ord vn) => PatBase f vn -> S.Set (IdentBase f vn)
patIdents :: forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents (Id vn
v f PatType
t SrcLoc
loc) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
vn -> f PatType -> SrcLoc -> IdentBase f vn
Ident vn
v f PatType
t SrcLoc
loc
patIdents (PatParens PatBase f vn
p SrcLoc
_) = forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents PatBase f vn
p
patIdents (TuplePat [PatBase f vn]
pats SrcLoc
_) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase f vn]
pats
patIdents (RecordPat [(Name, PatBase f vn)]
fs SrcLoc
_) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, PatBase f vn)]
fs
patIdents Wildcard {} = forall a. Monoid a => a
mempty
patIdents (PatAscription PatBase f vn
p TypeExp vn
_ SrcLoc
_) = forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents PatBase f vn
p
patIdents PatLit {} = forall a. Monoid a => a
mempty
patIdents (PatConstr Name
_ f PatType
_ [PatBase f vn]
ps SrcLoc
_) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase f vn]
ps
patIdents (PatAttr AttrInfo vn
_ PatBase f vn
p SrcLoc
_) = forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents PatBase f vn
p

-- | The set of names bound in a pattern.
patNames :: (Functor f, Ord vn) => PatBase f vn -> S.Set vn
patNames :: forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames (Id vn
v f PatType
_ SrcLoc
_) = forall a. a -> Set a
S.singleton vn
v
patNames (PatParens PatBase f vn
p SrcLoc
_) = forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase f vn
p
patNames (TuplePat [PatBase f vn]
pats SrcLoc
_) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase f vn]
pats
patNames (RecordPat [(Name, PatBase f vn)]
fs SrcLoc
_) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, PatBase f vn)]
fs
patNames Wildcard {} = forall a. Monoid a => a
mempty
patNames (PatAscription PatBase f vn
p TypeExp vn
_ SrcLoc
_) = forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase f vn
p
patNames PatLit {} = forall a. Monoid a => a
mempty
patNames (PatConstr Name
_ f PatType
_ [PatBase f vn]
ps SrcLoc
_) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase f vn]
ps
patNames (PatAttr AttrInfo vn
_ PatBase f vn
p SrcLoc
_) = forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase f vn
p

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

-- | The type of values bound by the pattern.
patternType :: PatBase Info VName -> PatType
patternType :: PatBase Info VName -> PatType
patternType (Wildcard (Info PatType
t) SrcLoc
_) = PatType
t
patternType (PatParens PatBase Info VName
p SrcLoc
_) = PatBase Info VName -> PatType
patternType PatBase Info VName
p
patternType (Id VName
_ (Info PatType
t) SrcLoc
_) = PatType
t
patternType (TuplePat [PatBase Info VName]
pats SrcLoc
_) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map PatBase Info VName -> PatType
patternType [PatBase Info VName]
pats
patternType (RecordPat [(Name, PatBase Info VName)]
fs SrcLoc
_) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ PatBase Info VName -> PatType
patternType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatBase Info VName)]
fs
patternType (PatAscription PatBase Info VName
p TypeExp VName
_ SrcLoc
_) = PatBase Info VName -> PatType
patternType PatBase Info VName
p
patternType (PatLit PatLit
_ (Info PatType
t) SrcLoc
_) = PatType
t
patternType (PatConstr Name
_ (Info PatType
t) [PatBase Info VName]
_ SrcLoc
_) = PatType
t
patternType (PatAttr AttrInfo VName
_ PatBase Info VName
p SrcLoc
_) = PatBase Info VName -> PatType
patternType PatBase Info VName
p

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

-- | When viewed as a function parameter, does this pattern correspond
-- to a named parameter of some type?
patternParam :: PatBase Info VName -> (PName, StructType)
patternParam :: PatBase Info VName -> (PName, StructType)
patternParam (PatParens PatBase Info VName
p SrcLoc
_) =
  PatBase Info VName -> (PName, StructType)
patternParam PatBase Info VName
p
patternParam (PatAttr AttrInfo VName
_ PatBase Info VName
p SrcLoc
_) =
  PatBase Info VName -> (PName, StructType)
patternParam PatBase Info VName
p
patternParam (PatAscription (Id VName
v (Info PatType
t) SrcLoc
_) TypeExp VName
_ SrcLoc
_) =
  (VName -> PName
Named VName
v, forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)
patternParam (Id VName
v (Info PatType
t) SrcLoc
_) =
  (VName -> PName
Named VName
v, forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)
patternParam PatBase Info VName
p =
  (PName
Unnamed, PatBase Info VName -> StructType
patternStructType PatBase 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 =
  forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    [ (FilePath -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> FilePath
prettyString PrimType
t, PrimType
t)
      | PrimType
t <-
          PrimType
Bool
            forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
            forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
            forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
    ]

-- | The nature of something predefined.  For functions, 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] (RetTypeBase Size ())
  | IntrinsicType Liftedness [TypeParamBase VName] StructType
  | IntrinsicEquality -- Special cased.

intrinsicAcc :: (VName, Intrinsic)
intrinsicAcc :: (VName, Intrinsic)
intrinsicAcc =
  ( VName
acc_v,
    Liftedness -> [TypeParamBase VName] -> StructType -> Intrinsic
IntrinsicType Liftedness
SizeLifted [forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
t_v forall a. Monoid a => a
mempty] forall a b. (a -> b) -> a -> b
$
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$
        forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (forall v. v -> QualName v
qualName VName
acc_v) [forall {dim}. TypeArg dim
arg]
  )
  where
    acc_v :: VName
acc_v = Name -> Int -> VName
VName Name
"acc" Int
10
    t_v :: VName
t_v = Name -> Int -> VName
VName Name
"t" Int
11
    arg :: TypeArg dim
arg = forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (forall v. v -> QualName v
qualName VName
t_v) [])) forall a. Monoid a => a
mempty

-- | A map of all built-ins.
intrinsics :: M.Map VName Intrinsic
intrinsics :: Map VName Intrinsic
intrinsics =
  (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, Intrinsic)
intrinsicAcc] <>) forall a b. (a -> b) -> a -> b
$
    forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$
      forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {b}. Int -> (FilePath, b) -> (VName, b)
namify [Int
20 ..] forall a b. (a -> b) -> a -> b
$
        forall a b. (a -> b) -> [a] -> [b]
map forall {a} {c}. (a, ([PrimType], PrimType, c)) -> (a, Intrinsic)
primFun (forall k a. Map k a -> [(k, a)]
M.toList Map FilePath ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
Primitive.primFuns)
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map UnOp -> (FilePath, Intrinsic)
unOpFun [UnOp]
Primitive.allUnOps
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map BinOp -> (FilePath, Intrinsic)
binOpFun [BinOp]
Primitive.allBinOps
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map CmpOp -> (FilePath, Intrinsic)
cmpOpFun [CmpOp]
Primitive.allCmpOps
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map ConvOp -> (FilePath, Intrinsic)
convOpFun [ConvOp]
Primitive.allConvOps
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map IntType -> (FilePath, Intrinsic)
signFun [IntType]
Primitive.allIntTypes
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map IntType -> (FilePath, Intrinsic)
unsignFun [IntType]
Primitive.allIntTypes
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map
            PrimType -> (FilePath, Intrinsic)
intrinsicPrim
            ( forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
                forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
                forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
                forall a. [a] -> [a] -> [a]
++ [PrimType
Bool]
            )
          forall a. [a] -> [a] -> [a]
++
          -- This overrides the ! from Primitive.
          [ ( FilePath
"!",
              [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> Intrinsic
IntrinsicOverloadedFun
                ( forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
                    forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
                    forall a. [a] -> [a] -> [a]
++ [PrimType
Bool]
                )
                [forall a. Maybe a
Nothing]
                forall a. Maybe a
Nothing
            )
          ]
          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.
          forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe BinOp -> Maybe (FilePath, Intrinsic)
mkIntrinsicBinOp [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
          forall a. [a] -> [a] -> [a]
++ [ ( FilePath
"flatten",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                   [forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
n, VName
m]) forall {dim}. ScalarTypeBase dim ()
t_a]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
k]
                   forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
k]) forall {dim}. ScalarTypeBase dim ()
t_a
               ),
               ( FilePath
"unflatten",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
n]) forall {dim}. ScalarTypeBase dim ()
t_a
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
k, VName
m]
                   forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
k, VName
m]) forall {dim}. ScalarTypeBase dim ()
t_a
               ),
               ( FilePath
"concat",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                   [forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n], forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m]]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
k]
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k]
               ),
               ( FilePath
"rotate",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                   [forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64, forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
               ),
               ( FilePath
"transpose",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                   [forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n, VName
m]]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
n]
               ),
               ( FilePath
"scatter",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_l]
                   [ forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique ([VName] -> Shape Size
shape [VName
n]) forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
l]) (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
l]) forall {dim}. ScalarTypeBase dim ()
t_a
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique ([VName] -> Shape Size
shape [VName
n]) forall {dim}. ScalarTypeBase dim ()
t_a
               ),
               ( FilePath
"scatter_2d",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_l]
                   [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n, VName
m],
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
l]) (forall {dim} {as}. Int -> ScalarTypeBase dim as
tupInt64 Int
2),
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
l]) forall {dim}. ScalarTypeBase dim ()
t_a
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n, VName
m]
               ),
               ( FilePath
"scatter_3d",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l]
                   [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n, VName
m, VName
k],
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
l]) (forall {dim} {as}. Int -> ScalarTypeBase dim as
tupInt64 Int
3),
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
l]) forall {dim}. ScalarTypeBase dim ()
t_a
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n, VName
m, VName
k]
               ),
               ( FilePath
"zip",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                   [forall {dim}. Shape dim -> TypeBase dim ()
arr_a ([VName] -> Shape Size
shape [VName
n]), forall {dim}. Shape dim -> TypeBase dim ()
arr_b ([VName] -> Shape Size
shape [VName
n])]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}.
TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
tuple_uarr (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a) (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b)
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
               ),
               ( FilePath
"unzip",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                   [forall {dim}.
TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
tuple_arr (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a) (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b) forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                   forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n], forall {dim}. Shape dim -> TypeBase dim ()
arr_b forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]]
               ),
               ( FilePath
"hist_1d",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                     forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m],
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
n]) (forall {dim} {as}. Int -> ScalarTypeBase dim as
tupInt64 Int
1),
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a ([VName] -> Shape Size
shape [VName
n])
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m]
               ),
               ( FilePath
"hist_2d",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                     forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k],
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
n]) (forall {dim} {as}. Int -> ScalarTypeBase dim as
tupInt64 Int
2),
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a ([VName] -> Shape Size
shape [VName
n])
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k]
               ),
               ( FilePath
"hist_3d",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                     forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k, VName
l],
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([VName] -> Shape Size
shape [VName
n]) (forall {dim} {as}. Int -> ScalarTypeBase dim as
tupInt64 Int
3),
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a ([VName] -> Shape Size
shape [VName
n])
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k, VName
l]
               ),
               ( FilePath
"map",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b,
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_b
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
               ),
               ( FilePath
"reduce",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a
               ),
               ( FilePath
"reduce_comm",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a
               ),
               ( FilePath
"scan",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
               ),
               ( FilePath
"partition",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int32),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   ( forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$
                       forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord
                         [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k],
                           forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique ([VName] -> Shape Size
shape [VName
n]) (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                         ]
                   )
               ),
               ( FilePath
"acc_write",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
sp_k, TypeParamBase VName
tp_a]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall {dim}. TypeBase dim () -> ScalarTypeBase dim ()
accType StructType
arr_ka,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. TypeBase dim () -> ScalarTypeBase dim ()
accType StructType
arr_ka
               ),
               ( FilePath
"scatter_stream",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_k, TypeParamBase VName
sp_n]
                   [ StructType
uarr_ka,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall {dim}. TypeBase dim () -> ScalarTypeBase dim ()
accType StructType
arr_ka)
                       forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` ( forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b
                                 forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall {dim}. TypeBase dim () -> ScalarTypeBase dim ()
accType forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k])
                             ),
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_b forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] StructType
uarr_ka
               ),
               ( FilePath
"hist_stream",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_k, TypeParamBase VName
sp_n]
                   [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k],
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a),
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall {dim}. TypeBase dim () -> ScalarTypeBase dim ()
accType StructType
arr_ka)
                       forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` ( forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b
                                 forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall {dim}. TypeBase dim () -> ScalarTypeBase dim ()
accType forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k])
                             ),
                     forall {dim}. Shape dim -> TypeBase dim ()
arr_b forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                   forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k]
               ),
               ( FilePath
"jvp2",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord [forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b, forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b]
               ),
               ( FilePath
"vjp2",
                 [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                   [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                   [ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a forall {as} {dim}.
Monoid as =>
TypeBase dim () -> TypeBase dim as -> TypeBase dim as
`arr` forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a,
                     forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b
                   ]
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                   forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar
                   forall a b. (a -> b) -> a -> b
$ forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord [forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_b, forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall {dim}. ScalarTypeBase dim ()
t_a]
               )
             ]
          forall a. [a] -> [a] -> [a]
++
          -- Experimental LMAD ones.
          [ ( FilePath
"flat_index_2d",
              [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                [ forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n],
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                ]
                forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m, VName
k]
                forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a
                forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k]
            ),
            ( FilePath
"flat_update_2d",
              [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l]
                [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n],
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k, VName
l]
                ]
                forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
            ),
            ( FilePath
"flat_index_3d",
              [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                [ forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n],
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                ]
                forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m, VName
k, VName
l]
                forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a
                forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k, VName
l]
            ),
            ( FilePath
"flat_update_3d",
              [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l, TypeParamBase VName
sp_p]
                [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n],
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k, VName
l, VName
p]
                ]
                forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
            ),
            ( FilePath
"flat_index_4d",
              [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                [ forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n],
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                ]
                forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m, VName
k, VName
l, VName
p]
                forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
arr_a
                forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
m, VName
k, VName
l, VName
p]
            ),
            ( FilePath
"flat_update_4d",
              [TypeParamBase VName] -> [StructType] -> StructRetType -> Intrinsic
IntrinsicPolyFun
                [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l, TypeParamBase VName
sp_p, TypeParamBase VName
sp_q]
                [ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n],
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                  forall {dim}. Shape dim -> TypeBase dim ()
arr_a forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
k, VName
l, VName
p, VName
q]
                ]
                forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                forall a b. (a -> b) -> a -> b
$ forall {dim}. Shape dim -> TypeBase dim ()
uarr_a
                forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Size
shape [VName
n]
            )
          ]
  where
    [VName
a, VName
b, VName
n, VName
m, VName
k, VName
l, VName
p, VName
q] = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Int -> VName
VName (forall a b. (a -> b) -> [a] -> [b]
map FilePath -> Name
nameFromString [FilePath
"a", FilePath
"b", FilePath
"n", FilePath
"m", FilePath
"k", FilePath
"l", FilePath
"p", FilePath
"q"]) [Int
0 ..]

    t_a :: ScalarTypeBase dim ()
t_a = forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (forall v. v -> QualName v
qualName VName
a) []
    arr_a :: Shape dim -> TypeBase dim ()
arr_a Shape dim
s = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique Shape dim
s forall {dim}. ScalarTypeBase dim ()
t_a
    uarr_a :: Shape dim -> TypeBase dim ()
uarr_a Shape dim
s = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique Shape dim
s forall {dim}. ScalarTypeBase dim ()
t_a
    tp_a :: TypeParamBase VName
tp_a = forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
a forall a. Monoid a => a
mempty

    t_b :: ScalarTypeBase dim ()
t_b = forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (forall v. v -> QualName v
qualName VName
b) []
    arr_b :: Shape dim -> TypeBase dim ()
arr_b Shape dim
s = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique Shape dim
s forall {dim}. ScalarTypeBase dim ()
t_b
    uarr_b :: Shape dim -> TypeBase dim ()
uarr_b Shape dim
s = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique Shape dim
s forall {dim}. ScalarTypeBase dim ()
t_b
    tp_b :: TypeParamBase VName
tp_b = forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
b forall a. Monoid a => a
mempty

    [TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l, TypeParamBase VName
sp_p, TypeParamBase VName
sp_q] = forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) [VName
n, VName
m, VName
k, VName
l, VName
p, VName
q]

    shape :: [VName] -> Shape Size
shape = forall dim. [dim] -> Shape dim
Shape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (QualName VName -> Size
NamedSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName)

    tuple_arr :: TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
tuple_arr TypeBase dim ()
x TypeBase dim ()
y Shape dim
s =
      forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array
        ()
        Uniqueness
Nonunique
        Shape dim
s
        (forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [TypeBase dim ()
x, TypeBase dim ()
y]))
    tuple_uarr :: TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
tuple_uarr TypeBase dim ()
x TypeBase dim ()
y Shape dim
s = forall {dim}.
TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
tuple_arr TypeBase dim ()
x TypeBase dim ()
y Shape dim
s forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Unique

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

    arr_ka :: StructType
arr_ka = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique (forall dim. [dim] -> Shape dim
Shape [QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
k]) forall {dim}. ScalarTypeBase dim ()
t_a
    uarr_ka :: StructType
uarr_ka = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique (forall dim. [dim] -> Shape dim
Shape [QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
k]) forall {dim}. ScalarTypeBase dim ()
t_a

    accType :: TypeBase dim () -> ScalarTypeBase dim ()
accType TypeBase dim ()
t =
      forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Unique (forall v. v -> QualName v
qualName (forall a b. (a, b) -> a
fst (VName, Intrinsic)
intrinsicAcc)) [forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType TypeBase dim ()
t forall a. Monoid a => a
mempty]

    namify :: Int -> (FilePath, b) -> (VName, b)
namify Int
i (FilePath
x, b
y) = (Name -> Int -> VName
VName (FilePath -> Name
nameFromString FilePath
x) Int
i, b
y)

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

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

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

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

    convOpFun :: ConvOp -> (FilePath, Intrinsic)
convOpFun ConvOp
cop = (forall a. Pretty a => a -> FilePath
prettyString ConvOp
cop, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [PrimType -> PrimType
unPrim PrimType
ft] 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 -> (FilePath, Intrinsic)
signFun IntType
t = (FilePath
"sign_" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
prettyString IntType
t, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [IntType -> PrimType
Unsigned IntType
t] forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
t)

    unsignFun :: IntType -> (FilePath, Intrinsic)
unsignFun IntType
t = (FilePath
"unsign_" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
prettyString IntType
t, [PrimType] -> PrimType -> Intrinsic
IntrinsicMonoFun [IntType -> PrimType
Signed IntType
t] 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.Unit = PrimType
Bool

    intrinsicPrim :: PrimType -> (FilePath, Intrinsic)
intrinsicPrim PrimType
t = (forall a. Pretty a => a -> FilePath
prettyString PrimType
t, Liftedness -> [TypeParamBase VName] -> StructType -> Intrinsic
IntrinsicType Liftedness
Unlifted [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t)

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

    mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic)
    mkIntrinsicBinOp :: BinOp -> Maybe (FilePath, Intrinsic)
mkIntrinsicBinOp BinOp
op = do
      Intrinsic
op' <- BinOp -> Maybe Intrinsic
intrinsicBinOp BinOp
op
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Pretty a => a -> FilePath
prettyString BinOp
op, Intrinsic
op')

    binOp :: [PrimType] -> Maybe Intrinsic
binOp [PrimType]
ts = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> Intrinsic
IntrinsicOverloadedFun [PrimType]
ts [forall a. Maybe a
Nothing, forall a. Maybe a
Nothing] forall a. Maybe a
Nothing
    ordering :: Maybe Intrinsic
ordering = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> Intrinsic
IntrinsicOverloadedFun [PrimType]
anyPrimType [forall a. Maybe a
Nothing, forall a. Maybe a
Nothing] (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 = forall a. a -> Maybe a
Just Intrinsic
IntrinsicEquality
    intrinsicBinOp BinOp
NotEqual = 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
_ = forall a. Maybe a
Nothing

    tupInt64 :: Int -> ScalarTypeBase dim as
tupInt64 Int
1 =
      forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    tupInt64 Int
x =
      forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate Int
x forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- | Is this include part of the built-in prelude?
isBuiltin :: FilePath -> Bool
isBuiltin :: FilePath -> Bool
isBuiltin = (forall a. Eq a => a -> a -> Bool
== FilePath
"/prelude") forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
takeDirectory

-- | Is the position of this thing builtin as per 'isBuiltin'?  Things
-- without location are considered not built-in.
isBuiltinLoc :: Located a => a -> Bool
isBuiltinLoc :: forall a. Located a => a -> Bool
isBuiltinLoc a
x =
  case forall a. Located a => a -> Loc
locOf a
x of
    Loc
NoLoc -> Bool
False
    Loc Pos
pos Pos
_ -> FilePath -> Bool
isBuiltin forall a b. (a -> b) -> a -> b
$ Pos -> FilePath
posFile Pos
pos

-- | 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 = forall a (f :: * -> *). (Num a, Ord a, Foldable f) => f a -> a
maxinum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> Int
baseTag forall a b. (a -> b) -> a -> b
$ 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 :: forall v. v -> QualName v
qualName = forall vn. [vn] -> vn -> QualName vn
QualName []

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

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

-- | The modules imported by a single declaration.
decImports :: DecBase f vn -> [(String, SrcLoc)]
decImports :: forall (f :: * -> *) vn. DecBase f vn -> [(FilePath, SrcLoc)]
decImports (OpenDec ModExpBase f vn
x SrcLoc
_) = forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, SrcLoc)]
modExpImports ModExpBase f vn
x
decImports (ModDec ModBindBase f vn
md) = forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, SrcLoc)]
modExpImports forall a b. (a -> b) -> a -> b
$ 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
_) = forall (f :: * -> *) vn. DecBase f vn -> [(FilePath, SrcLoc)]
decImports DecBase f vn
d
decImports (ImportDec FilePath
x f FilePath
_ SrcLoc
loc) = [(FilePath
x, SrcLoc
loc)]

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

-- | The set of module types used in any exported (non-local)
-- declaration.
progModuleTypes :: ProgBase Info VName -> S.Set VName
progModuleTypes :: ProgBase Info VName -> Set VName
progModuleTypes ProgBase Info VName
prog = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> Set VName
reach Set VName
mtypes_used
  where
    -- Fixed point iteration.
    reach :: VName -> Set VName
reach VName
v = forall a. a -> Set a
S.singleton VName
v forall a. Semigroup a => a -> a -> a
<> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> Set VName
reach) (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (Set VName)
reachable_from_mtype)

    reachable_from_mtype :: Map VName (Set VName)
reachable_from_mtype = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall {a} {f :: * -> *}. Ord a => DecBase f a -> Map a (Set a)
onDec forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs ProgBase Info VName
prog
      where
        onDec :: DecBase f a -> Map a (Set a)
onDec OpenDec {} = forall a. Monoid a => a
mempty
        onDec ModDec {} = forall a. Monoid a => a
mempty
        onDec (SigDec SigBindBase f a
sb) =
          forall k a. k -> a -> Map k a
M.singleton (forall (f :: * -> *) vn. SigBindBase f vn -> vn
sigName SigBindBase f a
sb) (forall {a} {f :: * -> *}. Ord a => SigExpBase f a -> Set a
onSigExp (forall (f :: * -> *) vn. SigBindBase f vn -> SigExpBase f vn
sigExp SigBindBase f a
sb))
        onDec TypeDec {} = forall a. Monoid a => a
mempty
        onDec ValDec {} = forall a. Monoid a => a
mempty
        onDec (LocalDec DecBase f a
d SrcLoc
_) = DecBase f a -> Map a (Set a)
onDec DecBase f a
d
        onDec ImportDec {} = forall a. Monoid a => a
mempty

        onSigExp :: SigExpBase f a -> Set a
onSigExp (SigVar QualName a
v f (Map VName VName)
_ SrcLoc
_) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ 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 [SpecBase f a]
ss SrcLoc
_) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap SpecBase f a -> Set a
onSpec [SpecBase f a]
ss
        onSigExp (SigWith SigExpBase f a
e TypeRefBase 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 forall a. Semigroup a => a -> a -> a
<> SigExpBase f a -> Set a
onSigExp SigExpBase f a
e2

        onSpec :: SpecBase f a -> Set a
onSpec ValSpec {} = forall a. Monoid a => a
mempty
        onSpec TypeSpec {} = forall a. Monoid a => a
mempty
        onSpec TypeAbbrSpec {} = forall a. Monoid a => a
mempty
        onSpec (ModSpec a
vn SigExpBase f a
e Maybe DocComment
_ SrcLoc
_) = forall a. a -> Set a
S.singleton a
vn forall a. Semigroup a => a -> a -> a
<> SigExpBase f a -> Set a
onSigExp SigExpBase f a
e
        onSpec (IncludeSpec SigExpBase f a
e SrcLoc
_) = SigExpBase f a -> Set a
onSigExp SigExpBase f a
e

    mtypes_used :: Set VName
mtypes_used = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall {f :: * -> *}. DecBase f VName -> Set VName
onDec forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs ProgBase Info VName
prog
      where
        onDec :: DecBase f VName -> Set VName
onDec (OpenDec ModExpBase f VName
x SrcLoc
_) = ModExpBase f VName -> Set VName
onModExp ModExpBase f VName
x
        onDec (ModDec ModBindBase f VName
md) =
          forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall {a} {f :: * -> *}. Ord a => SigExpBase f a -> Set a
onSigExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) (forall (f :: * -> *) vn.
ModBindBase f vn -> Maybe (SigExpBase f vn, f (Map VName VName))
modSignature ModBindBase f VName
md) forall a. Semigroup a => a -> a -> a
<> ModExpBase f VName -> Set VName
onModExp (forall (f :: * -> *) vn. ModBindBase f vn -> ModExpBase f vn
modExp ModBindBase f VName
md)
        onDec SigDec {} = forall a. Monoid a => a
mempty
        onDec TypeDec {} = forall a. Monoid a => a
mempty
        onDec ValDec {} = forall a. Monoid a => a
mempty
        onDec LocalDec {} = forall a. Monoid a => a
mempty
        onDec ImportDec {} = forall a. Monoid a => a
mempty

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

        onModParam :: ModParamBase f VName -> Set VName
onModParam = forall {a} {f :: * -> *}. Ord a => SigExpBase f a -> Set a
onSigExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
_) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ 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 {} = forall a. Monoid a => a
mempty
        onSigExp (SigWith SigExpBase f a
e TypeRefBase 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 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 :: FilePath -> Maybe ((FilePath, FilePath, Maybe FilePath), FilePath)
identifierReference (Char
'`' : FilePath
s)
  | (FilePath
identifier, Char
'`' : Char
'@' : FilePath
s') <- forall a. (a -> Bool) -> [a] -> ([a], [a])
break (forall a. Eq a => a -> a -> Bool
== Char
'`') FilePath
s,
    (FilePath
namespace, FilePath
s'') <- forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isAlpha FilePath
s',
    Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Bool
null FilePath
namespace =
      case FilePath
s'' of
        Char
'@' : Char
'"' : FilePath
s'''
          | (FilePath
file, Char
'"' : FilePath
s'''') <- forall a. (a -> Bool) -> [a] -> ([a], [a])
span (forall a. Eq a => a -> a -> Bool
/= Char
'"') FilePath
s''' ->
              forall a. a -> Maybe a
Just ((FilePath
identifier, FilePath
namespace, forall a. a -> Maybe a
Just FilePath
file), FilePath
s'''')
        FilePath
_ -> forall a. a -> Maybe a
Just ((FilePath
identifier, FilePath
namespace, forall a. Maybe a
Nothing), FilePath
s'')
identifierReference FilePath
_ = 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 =
  forall b a. b -> (a -> b) -> Maybe a -> b
maybe BinOp
Backtick forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
s') forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$
      forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (forall a. a -> Down a
Down forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$
        forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> FilePath
prettyString [BinOp]
operators) [BinOp]
operators
  where
    s' :: FilePath
s' = Name -> FilePath
nameToString Name
s
    operators :: [BinOp]
    operators :: [BinOp]
operators = [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound :: BinOp]

-- | Find instances of typed holes in the program.
progHoles :: ProgBase Info VName -> [(Loc, StructType)]
progHoles :: ProgBase Info VName -> [(Loc, StructType)]
progHoles = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase Info VName -> [(Loc, StructType)]
holesInDec forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs
  where
    holesInDec :: DecBase Info VName -> [(Loc, StructType)]
holesInDec (ValDec ValBindBase Info VName
vb) = ExpBase Info VName -> [(Loc, StructType)]
holesInExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBindBase Info VName
vb
    holesInDec (ModDec ModBindBase Info VName
me) = ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ModBindBase f vn -> ModExpBase f vn
modExp ModBindBase Info VName
me
    holesInDec (OpenDec ModExpBase Info VName
me SrcLoc
_) = ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp ModExpBase Info VName
me
    holesInDec (LocalDec DecBase Info VName
d SrcLoc
_) = DecBase Info VName -> [(Loc, StructType)]
holesInDec DecBase Info VName
d
    holesInDec TypeDec {} = forall a. Monoid a => a
mempty
    holesInDec SigDec {} = forall a. Monoid a => a
mempty
    holesInDec ImportDec {} = forall a. Monoid a => a
mempty

    holesInModExp :: ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp (ModDecs [DecBase Info VName]
ds SrcLoc
_) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase Info VName -> [(Loc, StructType)]
holesInDec [DecBase Info VName]
ds
    holesInModExp (ModParens ModExpBase Info VName
me SrcLoc
_) = ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp ModExpBase Info VName
me
    holesInModExp (ModApply ModExpBase Info VName
x ModExpBase Info VName
y Info (Map VName VName)
_ Info (Map VName VName)
_ SrcLoc
_) = ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp ModExpBase Info VName
x forall a. Semigroup a => a -> a -> a
<> ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp ModExpBase Info VName
y
    holesInModExp (ModAscript ModExpBase Info VName
me SigExpBase Info VName
_ Info (Map VName VName)
_ SrcLoc
_) = ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp ModExpBase Info VName
me
    holesInModExp (ModLambda ModParamBase Info VName
_ Maybe (SigExpBase Info VName, Info (Map VName VName))
_ ModExpBase Info VName
me SrcLoc
_) = ModExpBase Info VName -> [(Loc, StructType)]
holesInModExp ModExpBase Info VName
me
    holesInModExp ModVar {} = forall a. Monoid a => a
mempty
    holesInModExp ModImport {} = forall a. Monoid a => a
mempty

    holesInExp :: ExpBase Info VName -> [(Loc, StructType)]
holesInExp = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> s
execState forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {m :: * -> *}.
MonadState [(Loc, StructType)] m =>
ExpBase Info VName -> m (ExpBase Info VName)
onExp

    onExp :: ExpBase Info VName -> m (ExpBase Info VName)
onExp e :: ExpBase Info VName
e@(Hole (Info PatType
t) SrcLoc
loc) = do
      forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((forall a. Located a => a -> Loc
locOf SrcLoc
loc, forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) :)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase Info VName
e
    onExp ExpBase Info VName
e = forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: ExpBase Info VName -> m (ExpBase Info VName)
mapOnExp = ExpBase Info VName -> m (ExpBase Info VName)
onExp}) ExpBase Info VName
e

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

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

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

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

-- | A slice with no type annotations.
type UncheckedSlice = SliceBase 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 UncheckedPat = PatBase 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