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

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

    -- * 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 qualified Data.Map.Strict as M
import Data.Maybe
import Data.Ord
import qualified Data.Set as S
import Futhark.Util (maxinum)
import Futhark.Util.Pretty
import qualified Language.Futhark.Primitive as Primitive
import Language.Futhark.Syntax
import Language.Futhark.Traversals
import Language.Futhark.Tuple

-- | The name of the default program entry point (@main@).
defaultEntryPoint :: Name
defaultEntryPoint :: Name
defaultEntryPoint = [Char] -> Name
nameFromString [Char]
"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 = Shape dim -> Int
forall dim. Shape dim -> Int
shapeRank (Shape dim -> Int)
-> (TypeBase dim as -> Shape dim) -> TypeBase dim as -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase dim as -> Shape dim
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
_ = Shape dim
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 = (Size -> ()) -> TypeBase Size as -> TypeBase () as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Size -> ()) -> TypeBase Size as -> TypeBase () as)
-> (Size -> ()) -> TypeBase Size as -> TypeBase () as
forall a b. (a -> b) -> a -> b
$ () -> Size -> ()
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
(DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool) -> Eq DimPos
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DimPos -> DimPos -> Bool
$c/= :: DimPos -> DimPos -> Bool
== :: DimPos -> DimPos -> Bool
$c== :: DimPos -> DimPos -> Bool
Eq, Eq DimPos
Eq DimPos
-> (DimPos -> DimPos -> Ordering)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> Bool)
-> (DimPos -> DimPos -> DimPos)
-> (DimPos -> DimPos -> DimPos)
-> Ord DimPos
DimPos -> DimPos -> Bool
DimPos -> DimPos -> Ordering
DimPos -> DimPos -> DimPos
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DimPos -> DimPos -> DimPos
$cmin :: DimPos -> DimPos -> DimPos
max :: DimPos -> DimPos -> DimPos
$cmax :: DimPos -> DimPos -> DimPos
>= :: DimPos -> DimPos -> Bool
$c>= :: DimPos -> DimPos -> Bool
> :: DimPos -> DimPos -> Bool
$c> :: DimPos -> DimPos -> Bool
<= :: DimPos -> DimPos -> Bool
$c<= :: DimPos -> DimPos -> Bool
< :: DimPos -> DimPos -> Bool
$c< :: DimPos -> DimPos -> Bool
compare :: DimPos -> DimPos -> Ordering
$ccompare :: DimPos -> DimPos -> Ordering
Ord, Int -> DimPos -> ShowS
[DimPos] -> ShowS
DimPos -> [Char]
(Int -> DimPos -> ShowS)
-> (DimPos -> [Char]) -> ([DimPos] -> ShowS) -> Show DimPos
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [DimPos] -> ShowS
$cshowList :: [DimPos] -> ShowS
show :: DimPos -> [Char]
$cshow :: DimPos -> [Char]
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 = Set VName -> DimPos -> TypeBase fdim als -> f (TypeBase tdim als)
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
forall a. Monoid a => a
mempty DimPos
PosImmediate
  where
    go ::
      forall als'.
      S.Set VName ->
      DimPos ->
      TypeBase fdim als' ->
      f (TypeBase tdim als')
    go :: forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b t :: TypeBase fdim als'
t@Array {} =
      (fdim -> f tdim)
-> (als' -> f als') -> TypeBase fdim als' -> f (TypeBase tdim als')
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (Set VName -> DimPos -> fdim -> f tdim
f Set VName
bound DimPos
b) als' -> f als'
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase fdim als'
t
    go Set VName
bound DimPos
b (Scalar (Record Map Name (TypeBase fdim als')
fields)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> (Map Name (TypeBase tdim als') -> ScalarTypeBase tdim als')
-> Map Name (TypeBase tdim als')
-> TypeBase tdim als'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase tdim als') -> ScalarTypeBase tdim als'
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase tdim als') -> TypeBase tdim als')
-> f (Map Name (TypeBase tdim als')) -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase fdim als' -> f (TypeBase tdim als'))
-> Map Name (TypeBase fdim als')
-> f (Map Name (TypeBase tdim als'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b) Map Name (TypeBase fdim als')
fields
    go Set VName
bound DimPos
b (Scalar (TypeVar als'
as Uniqueness
u QualName VName
tn [TypeArg fdim]
targs)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> f (ScalarTypeBase tdim als') -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (als'
-> Uniqueness
-> QualName VName
-> [TypeArg tdim]
-> ScalarTypeBase tdim als'
forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar als'
as Uniqueness
u QualName VName
tn ([TypeArg tdim] -> ScalarTypeBase tdim als')
-> f [TypeArg tdim] -> f (ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg fdim -> f (TypeArg tdim))
-> [TypeArg fdim] -> f [TypeArg tdim]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Set VName -> DimPos -> TypeArg fdim -> f (TypeArg tdim)
onTypeArg Set VName
bound DimPos
b) [TypeArg fdim]
targs)
    go Set VName
bound DimPos
b (Scalar (Sum Map Name [TypeBase fdim als']
cs)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> (Map Name [TypeBase tdim als'] -> ScalarTypeBase tdim als')
-> Map Name [TypeBase tdim als']
-> TypeBase tdim als'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase tdim als'] -> ScalarTypeBase tdim als'
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase tdim als'] -> TypeBase tdim als')
-> f (Map Name [TypeBase tdim als']) -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase fdim als'] -> f [TypeBase tdim als'])
-> Map Name [TypeBase fdim als']
-> f (Map Name [TypeBase tdim als'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase fdim als' -> f (TypeBase tdim als'))
-> [TypeBase fdim als'] -> f [TypeBase tdim als']
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b)) Map Name [TypeBase fdim als']
cs
    go Set VName
_ DimPos
_ (Scalar (Prim PrimType
t)) =
      TypeBase tdim als' -> f (TypeBase tdim als')
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase tdim als' -> f (TypeBase tdim als'))
-> TypeBase tdim als' -> f (TypeBase tdim als')
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> ScalarTypeBase tdim als' -> TypeBase tdim als'
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase tdim als'
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    go Set VName
bound DimPos
_ (Scalar (Arrow als'
als PName
p TypeBase fdim ()
t1 (RetType [VName]
dims TypeBase fdim als'
t2))) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> f (ScalarTypeBase tdim als') -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (als'
-> PName
-> TypeBase tdim ()
-> RetTypeBase tdim als'
-> ScalarTypeBase tdim als'
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow als'
als PName
p (TypeBase tdim ()
 -> RetTypeBase tdim als' -> ScalarTypeBase tdim als')
-> f (TypeBase tdim ())
-> f (RetTypeBase tdim als' -> ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> DimPos -> TypeBase fdim () -> f (TypeBase tdim ())
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosParam TypeBase fdim ()
t1 f (RetTypeBase tdim als' -> ScalarTypeBase tdim als')
-> f (RetTypeBase tdim als') -> f (ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([VName] -> TypeBase tdim als' -> RetTypeBase tdim als'
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase tdim als' -> RetTypeBase tdim als')
-> f (TypeBase tdim als') -> f (RetTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosReturn TypeBase fdim als'
t2))
      where
        bound' :: Set VName
bound' =
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims
            Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> case PName
p of
              Named VName
p' -> VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert VName
p' Set VName
bound
              PName
Unnamed -> Set VName
bound

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

-- | 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)) = ([TypeBase shape as] -> Uniqueness)
-> [[TypeBase shape as]] -> Uniqueness
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((TypeBase shape as -> Uniqueness)
-> [TypeBase shape as] -> Uniqueness
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase shape as -> Uniqueness
forall shape as. TypeBase shape as -> Uniqueness
uniqueness) ([[TypeBase shape as]] -> Uniqueness)
-> [[TypeBase shape as]] -> Uniqueness
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase shape as] -> [[TypeBase shape as]]
forall k a. Map k a -> [a]
M.elems Map Name [TypeBase shape as]
ts
uniqueness (Scalar (Record Map Name (TypeBase shape as)
fs)) = (TypeBase shape as -> Uniqueness)
-> [TypeBase shape as] -> Uniqueness
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase shape as -> Uniqueness
forall shape as. TypeBase shape as -> Uniqueness
uniqueness ([TypeBase shape as] -> Uniqueness)
-> [TypeBase shape as] -> Uniqueness
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase shape as) -> [TypeBase shape as]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase shape as)
fs
uniqueness TypeBase shape as
_ = Uniqueness
Nonunique

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

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

-- | @diet t@ returns a description of how a function parameter of
-- type @t@ might consume its argument.
diet :: TypeBase shape as -> Diet
diet :: forall shape as. TypeBase shape as -> Diet
diet (Scalar (Record Map Name (TypeBase shape as)
ets)) = Map Name Diet -> Diet
RecordDiet (Map Name Diet -> Diet) -> Map Name Diet -> Diet
forall a b. (a -> b) -> a -> b
$ (TypeBase shape as -> Diet)
-> Map Name (TypeBase shape as) -> Map Name Diet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape as -> Diet
forall shape as. TypeBase shape as -> Diet
diet Map Name (TypeBase shape as)
ets
diet (Scalar (Prim PrimType
_)) = Diet
Observe
diet (Scalar (Arrow as
_ PName
_ TypeBase shape ()
t1 (RetType [VName]
_ TypeBase shape as
t2))) = Diet -> Diet -> Diet
FuncDiet (TypeBase shape () -> Diet
forall shape as. TypeBase shape as -> Diet
diet TypeBase shape ()
t1) (TypeBase shape as -> Diet
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 (Map Name [Diet] -> Diet) -> Map Name [Diet] -> Diet
forall a b. (a -> b) -> a -> b
$ ([TypeBase shape as] -> [Diet])
-> Map Name [TypeBase shape as] -> Map Name [Diet]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase shape as -> Diet) -> [TypeBase shape as] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase shape as -> Diet
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 = (TypeBase () as -> () -> TypeBase () ())
-> () -> TypeBase () as -> TypeBase () ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip TypeBase () as -> () -> TypeBase () ()
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases () (TypeBase () as -> TypeBase () ())
-> (TypeBase dim as -> TypeBase () as)
-> TypeBase dim as
-> TypeBase () ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (dim -> ()) -> TypeBase dim as -> TypeBase () as
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (() -> dim -> ()
forall a b. a -> b -> a
const ())

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

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

-- | @peelArray n t@ returns the type resulting from peeling the first
-- @n@ array dimensions from @t@.  Returns @Nothing@ if @t@ has less
-- than @n@ dimensions.
peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray :: 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)
  | Shape dim -> Int
forall dim. Shape dim -> Int
shapeRank Shape dim
shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n =
      TypeBase dim as -> Maybe (TypeBase dim as)
forall a. a -> Maybe a
Just (TypeBase dim as -> Maybe (TypeBase dim as))
-> TypeBase dim as -> Maybe (TypeBase dim as)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
t TypeBase dim () -> (() -> as) -> TypeBase dim as
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` as -> () -> as
forall a b. a -> b -> a
const as
als
  | Bool
otherwise =
      as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array as
als Uniqueness
u (Shape dim -> ScalarTypeBase dim () -> TypeBase dim as)
-> Maybe (Shape dim)
-> Maybe (ScalarTypeBase dim () -> TypeBase dim as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Shape dim -> Maybe (Shape dim)
forall dim. Int -> Shape dim -> Maybe (Shape dim)
stripDims Int
n Shape dim
shape Maybe (ScalarTypeBase dim () -> TypeBase dim as)
-> Maybe (ScalarTypeBase dim ()) -> Maybe (TypeBase dim as)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ScalarTypeBase dim () -> Maybe (ScalarTypeBase dim ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase dim ()
t
peelArray Int
_ TypeBase dim as
_ = Maybe (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 = as -> Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
forall as dim.
Monoid as =>
as -> Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOfWithAliases as
forall a. Monoid a => a
mempty Uniqueness
u Shape dim
s (TypeBase dim as
t TypeBase dim as -> Uniqueness -> TypeBase dim as
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) =
  as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array (as
as1 as -> as -> as
forall a. Semigroup a => a -> a -> a
<> as
as2) Uniqueness
u (Shape dim
shape2 Shape dim -> Shape dim -> Shape dim
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) =
  as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array as
as Uniqueness
u Shape dim
shape ((as -> ()) -> ScalarTypeBase dim as -> ScalarTypeBase dim ()
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (() -> as -> ()
forall a b. a -> b -> a
const ()) ScalarTypeBase dim as
t)

-- | @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' <- Int -> Shape dim -> Maybe (Shape dim)
forall dim. Int -> Shape dim -> Maybe (Shape dim)
stripDims Int
n Shape dim
shape =
      as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
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 =
      ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
et TypeBase dim () -> Uniqueness -> TypeBase dim ()
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u TypeBase dim () -> as -> TypeBase dim as
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` as
als
stripArray Int
_ TypeBase dim as
t = TypeBase dim as
t

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

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

-- | 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 = ((Name, a) -> Name) -> [(Name, a)] -> [(Name, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, a) -> Name
forall a b. (a, b) -> a
fst ([(Name, a)] -> [(Name, a)]) -> [(Name, a)] -> [(Name, a)]
forall a b. (a -> b) -> a -> b
$ Map Name a -> [(Name, a)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name a
cs

-- | Is this a 'TypeParamType'?
isTypeParam :: TypeParamBase vn -> Bool
isTypeParam :: 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 (Bool -> Bool)
-> (TypeParamBase vn -> Bool) -> TypeParamBase vn -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase vn -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam

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

-- | 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 = [VName] -> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
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) ->
          (TypeBase d1 () -> as' -> TypeBase d1 as')
-> as' -> TypeBase d1 () -> TypeBase d1 as'
forall a b c. (a -> b -> c) -> b -> a -> c
flip TypeBase d1 () -> as' -> TypeBase d1 as'
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
setAliases (as'
als1 as' -> as' -> as'
forall a. Semigroup a => a -> a -> a
<> as'
als2)
            (TypeBase d1 () -> TypeBase d1 as')
-> m (TypeBase d1 ()) -> m (TypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( Uniqueness -> Shape d1 -> TypeBase d1 () -> TypeBase d1 ()
forall as dim.
Monoid as =>
Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOf (Uniqueness -> Uniqueness -> Uniqueness
forall a. Ord a => a -> a -> a
min Uniqueness
u1 Uniqueness
u2)
                    (Shape d1 -> TypeBase d1 () -> TypeBase d1 ())
-> m (Shape d1) -> m (TypeBase d1 () -> TypeBase d1 ())
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
                    m (TypeBase d1 () -> TypeBase d1 ())
-> m (TypeBase d1 ()) -> m (TypeBase d1 ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [VName] -> TypeBase d1 () -> TypeBase d2 () -> m (TypeBase d1 ())
forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound (ScalarTypeBase d1 () -> TypeBase d1 ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d1 ()
et1) (ScalarTypeBase d2 () -> TypeBase d2 ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d2 ()
et2)
                )
        (Scalar (Record Map Name (TypeBase d1 as')
f1), Scalar (Record Map Name (TypeBase d2 as')
f2)) ->
          ScalarTypeBase d1 as' -> TypeBase d1 as'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase d1 as' -> TypeBase d1 as')
-> (Map Name (TypeBase d1 as') -> ScalarTypeBase d1 as')
-> Map Name (TypeBase d1 as')
-> TypeBase d1 as'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase d1 as') -> ScalarTypeBase d1 as'
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record
            (Map Name (TypeBase d1 as') -> TypeBase d1 as')
-> m (Map Name (TypeBase d1 as')) -> m (TypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((TypeBase d1 as', TypeBase d2 as') -> m (TypeBase d1 as'))
-> Map Name (TypeBase d1 as', TypeBase d2 as')
-> m (Map Name (TypeBase d1 as'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as'))
-> (TypeBase d1 as', TypeBase d2 as') -> m (TypeBase d1 as')
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ([VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound)) ((TypeBase d1 as'
 -> TypeBase d2 as' -> (TypeBase d1 as', TypeBase d2 as'))
-> Map Name (TypeBase d1 as')
-> Map Name (TypeBase d2 as')
-> Map Name (TypeBase d1 as', TypeBase d2 as')
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name (TypeBase d1 as')
f1 Map Name (TypeBase d2 as')
f2)
        (Scalar (Sum Map Name [TypeBase d1 as']
cs1), Scalar (Sum Map Name [TypeBase d2 as']
cs2)) ->
          ScalarTypeBase d1 as' -> TypeBase d1 as'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase d1 as' -> TypeBase d1 as')
-> (Map Name [TypeBase d1 as'] -> ScalarTypeBase d1 as')
-> Map Name [TypeBase d1 as']
-> TypeBase d1 as'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase d1 as'] -> ScalarTypeBase d1 as'
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum
            (Map Name [TypeBase d1 as'] -> TypeBase d1 as')
-> m (Map Name [TypeBase d1 as']) -> m (TypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(TypeBase d1 as', TypeBase d2 as')] -> m [TypeBase d1 as'])
-> Map Name [(TypeBase d1 as', TypeBase d2 as')]
-> m (Map Name [TypeBase d1 as'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
              (((TypeBase d1 as', TypeBase d2 as') -> m (TypeBase d1 as'))
-> [(TypeBase d1 as', TypeBase d2 as')] -> m [TypeBase d1 as']
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as'))
-> (TypeBase d1 as', TypeBase d2 as') -> m (TypeBase d1 as')
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ([VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound)))
              (([TypeBase d1 as']
 -> [TypeBase d2 as'] -> [(TypeBase d1 as', TypeBase d2 as')])
-> Map Name [TypeBase d1 as']
-> Map Name [TypeBase d2 as']
-> Map Name [(TypeBase d1 as', TypeBase d2 as')]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith [TypeBase d1 as']
-> [TypeBase d2 as'] -> [(TypeBase d1 as', TypeBase d2 as')]
forall a b. [a] -> [b] -> [(a, b)]
zip Map Name [TypeBase d1 as']
cs1 Map Name [TypeBase d2 as']
cs2)
        ( Scalar (Arrow as'
als1 PName
p1 TypeBase d1 ()
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' = (PName -> Maybe VName) -> [PName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe PName -> Maybe VName
maybePName [PName
p1, PName
p2] [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims1 [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims2 [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
bound
             in ScalarTypeBase d1 as' -> TypeBase d1 as'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar
                  (ScalarTypeBase d1 as' -> TypeBase d1 as')
-> m (ScalarTypeBase d1 as') -> m (TypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( as'
-> PName
-> TypeBase d1 ()
-> RetTypeBase d1 as'
-> ScalarTypeBase d1 as'
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow (as'
als1 as' -> as' -> as'
forall a. Semigroup a => a -> a -> a
<> as'
als2) PName
p1
                          (TypeBase d1 () -> RetTypeBase d1 as' -> ScalarTypeBase d1 as')
-> m (TypeBase d1 ())
-> m (RetTypeBase d1 as' -> ScalarTypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName] -> TypeBase d1 () -> TypeBase d2 () -> m (TypeBase d1 ())
forall as'.
Monoid as' =>
[VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
matchDims' [VName]
bound' TypeBase d1 ()
a1 TypeBase d2 ()
a2
                          m (RetTypeBase d1 as' -> ScalarTypeBase d1 as')
-> m (RetTypeBase d1 as') -> m (ScalarTypeBase d1 as')
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([VName] -> TypeBase d1 as' -> RetTypeBase d1 as'
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims1 (TypeBase d1 as' -> RetTypeBase d1 as')
-> m (TypeBase d1 as') -> m (RetTypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName]
-> TypeBase d1 as' -> TypeBase d2 as' -> m (TypeBase d1 as')
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)
          ) ->
            ScalarTypeBase d1 as' -> TypeBase d1 as'
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase d1 as' -> TypeBase d1 as')
-> ([TypeArg d1] -> ScalarTypeBase d1 as')
-> [TypeArg d1]
-> TypeBase d1 as'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. as'
-> Uniqueness
-> QualName VName
-> [TypeArg d1]
-> ScalarTypeBase d1 as'
forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar (as'
als1 as' -> as' -> as'
forall a. Semigroup a => a -> a -> a
<> as'
als2) Uniqueness
u QualName VName
v
              ([TypeArg d1] -> TypeBase d1 as')
-> m [TypeArg d1] -> m (TypeBase d1 as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg d1 -> TypeArg d2 -> m (TypeArg d1))
-> [TypeArg d1] -> [TypeArg d2] -> m [TypeArg d1]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ([VName] -> TypeArg d1 -> TypeArg d2 -> m (TypeArg d1)
matchTypeArg [VName]
bound) [TypeArg d1]
targs1 [TypeArg d2]
targs2
        (TypeBase d1 as', TypeBase d2 as')
_ -> TypeBase d1 as' -> m (TypeBase d1 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
_ = TypeArg d1 -> m (TypeArg d1)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeArg d1
ta
    matchTypeArg [VName]
bound (TypeArgDim d1
x SrcLoc
loc) (TypeArgDim d2
y SrcLoc
_) =
      d1 -> SrcLoc -> TypeArg d1
forall dim. dim -> SrcLoc -> TypeArg dim
TypeArgDim (d1 -> SrcLoc -> TypeArg d1) -> m d1 -> m (SrcLoc -> TypeArg d1)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName] -> d1 -> d2 -> m d1
onDims [VName]
bound d1
x d2
y m (SrcLoc -> TypeArg d1) -> m SrcLoc -> m (TypeArg d1)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    matchTypeArg [VName]
_ TypeArg d1
a TypeArg d2
_ = TypeArg d1 -> m (TypeArg d1)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeArg d1
a

    maybePName :: PName -> Maybe VName
maybePName (Named VName
v) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
    maybePName PName
Unnamed = Maybe VName
forall a. Maybe a
Nothing

    onShapes :: [VName] -> Shape d1 -> Shape d2 -> m (Shape d1)
onShapes [VName]
bound Shape d1
shape1 Shape d2
shape2 =
      [d1] -> Shape d1
forall dim. [dim] -> Shape dim
Shape ([d1] -> Shape d1) -> m [d1] -> m (Shape d1)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (d1 -> d2 -> m d1) -> [d1] -> [d2] -> m [d1]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ([VName] -> d1 -> d2 -> m d1
onDims [VName]
bound) (Shape d1 -> [d1]
forall dim. Shape dim -> [dim]
shapeDims Shape d1
shape1) (Shape d2 -> [d2]
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 =
  as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
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 =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
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 =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim as) -> ScalarTypeBase dim as)
-> Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ (TypeBase dim as -> TypeBase dim as)
-> Map Name (TypeBase dim as) -> Map Name (TypeBase dim as)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (TypeBase dim as -> Uniqueness -> TypeBase dim as
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u) Map Name (TypeBase dim as)
ets
setUniqueness (Scalar (Sum Map Name [TypeBase dim as]
ets)) Uniqueness
u =
  ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim as] -> ScalarTypeBase dim as)
-> Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim as] -> [TypeBase dim as])
-> Map Name [TypeBase dim as] -> Map Name [TypeBase dim as]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim as -> TypeBase dim as)
-> [TypeBase dim as] -> [TypeBase dim as]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase dim as -> Uniqueness -> TypeBase dim as
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
u)) Map Name [TypeBase dim as]
ets
setUniqueness TypeBase dim as
t Uniqueness
_ = TypeBase dim as
t

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

-- | @t \`addAliases\` f@ returns @t@, but with any already present
-- 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 = ((asf -> ast) -> TypeBase dim asf -> TypeBase dim ast)
-> TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
forall a b c. (a -> b -> c) -> b -> a -> c
flip (asf -> ast) -> TypeBase dim asf -> TypeBase dim ast
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second

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

floatValueType :: FloatValue -> FloatType
floatValueType :: FloatValue -> FloatType
floatValueType 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 (IntType -> PrimType) -> IntType -> PrimType
forall a b. (a -> b) -> a -> b
$ IntValue -> IntType
intValueType IntValue
v
primValueType (UnsignedValue IntValue
v) = IntType -> PrimType
Unsigned (IntType -> PrimType) -> IntType -> PrimType
forall a b. (a -> b) -> a -> b
$ IntValue -> IntType
intValueType IntValue
v
primValueType (FloatValue FloatValue
v) = FloatType -> PrimType
FloatType (FloatType -> PrimType) -> FloatType -> PrimType
forall a b. (a -> b) -> a -> b
$ FloatValue -> FloatType
floatValueType FloatValue
v
primValueType BoolValue {} = PrimType
Bool

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

-- | The 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
_) = ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase Size Aliasing)
-> PrimType -> ScalarTypeBase Size Aliasing
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
_) = ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ [PatType] -> ScalarTypeBase Size Aliasing
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord ([PatType] -> ScalarTypeBase Size Aliasing)
-> [PatType] -> ScalarTypeBase Size Aliasing
forall a b. (a -> b) -> a -> b
$ (ExpBase Info VName -> PatType)
-> [ExpBase Info VName] -> [PatType]
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.
  ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Map Name PatType -> ScalarTypeBase Size Aliasing
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name PatType -> ScalarTypeBase Size Aliasing)
-> Map Name PatType -> ScalarTypeBase Size Aliasing
forall a b. (a -> b) -> a -> b
$ [Map Name PatType] -> Map Name PatType
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name PatType] -> Map Name PatType)
-> [Map Name PatType] -> Map Name PatType
forall a b. (a -> b) -> a -> b
$ [Map Name PatType] -> [Map Name PatType]
forall a. [a] -> [a]
reverse ([Map Name PatType] -> [Map Name PatType])
-> [Map Name PatType] -> [Map Name PatType]
forall a b. (a -> b) -> a -> b
$ (FieldBase Info VName -> Map Name PatType)
-> [FieldBase Info VName] -> [Map Name PatType]
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
_) = Name -> PatType -> Map Name PatType
forall k a. k -> a -> Map k a
M.singleton Name
name (PatType -> Map Name PatType) -> PatType -> Map Name PatType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatType
typeOf ExpBase Info VName
e
    record (RecordFieldImplicit VName
name (Info PatType
t) SrcLoc
_) =
      Name -> PatType -> Map Name PatType
forall k a. k -> a -> Map k a
M.singleton (VName -> Name
baseName VName
name) (PatType -> Map Name PatType) -> PatType -> Map Name PatType
forall a b. (a -> b) -> a -> b
$
        PatType
t
          PatType -> (Aliasing -> Aliasing) -> PatType
forall dim asf ast.
TypeBase dim asf -> (asf -> ast) -> TypeBase dim ast
`addAliases` Alias -> Aliasing -> Aliasing
forall a. Ord a => a -> Set a -> Set a
S.insert (VName -> Alias
AliasBound VName
name)
typeOf (ArrayLit [ExpBase Info VName]
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (StringLit [Word8]
vs SrcLoc
_) =
  Aliasing
-> Uniqueness -> Shape Size -> ScalarTypeBase Size () -> PatType
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array
    Aliasing
forall a. Monoid a => a
mempty
    Uniqueness
Unique
    ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [Int -> Size
ConstSize (Int -> Size) -> Int -> Size
forall a b. (a -> b) -> a -> b
$ [Word8] -> Int
forall i a. Num i => [a] -> i
genericLength [Word8]
vs])
    (PrimType -> ScalarTypeBase Size ()
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 PatType -> Aliasing -> PatType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` Aliasing
forall a. Monoid a => a
mempty
typeOf (RecordUpdate ExpBase Info VName
_ [Name]
_ ExpBase Info VName
_ (Info PatType
t) SrcLoc
_) = PatType
t
typeOf (Assert ExpBase Info VName
_ ExpBase Info VName
e Info [Char]
_ 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 [] TypeBase Size ()
t' = (PatBase Info VName -> StructRetType -> StructRetType)
-> StructRetType -> [PatBase Info VName] -> StructRetType
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((PName, TypeBase Size ()) -> StructRetType -> StructRetType
forall {dim}.
(PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow ((PName, TypeBase Size ()) -> StructRetType -> StructRetType)
-> (PatBase Info VName -> (PName, TypeBase Size ()))
-> PatBase Info VName
-> StructRetType
-> StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> (PName, TypeBase Size ())
patternParam) StructRetType
t [PatBase Info VName]
params
   in TypeBase Size ()
t' TypeBase Size () -> Aliasing -> PatType
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) =
      [VName] -> TypeBase dim () -> RetTypeBase dim ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim () -> RetTypeBase dim ())
-> TypeBase dim () -> RetTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> TypeBase dim ()
-> RetTypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () (VName -> PName
Named VName
v) TypeBase dim ()
x (RetTypeBase dim () -> ScalarTypeBase dim ())
-> RetTypeBase dim () -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase dim () -> RetTypeBase dim ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
dims) TypeBase dim ()
y
    arrow (PName
pn, TypeBase dim ()
tx) RetTypeBase dim ()
y =
      [VName] -> TypeBase dim () -> RetTypeBase dim ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim () -> RetTypeBase dim ())
-> TypeBase dim () -> RetTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> TypeBase dim ()
-> RetTypeBase dim ()
-> ScalarTypeBase dim ()
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, TypeBase Size (), Maybe VName)
_, Info (PName
pn, TypeBase Size ()
pt2)) (Info PatRetType
ret, Info [VName]
_) SrcLoc
_) =
  ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Aliasing
-> PName
-> TypeBase Size ()
-> PatRetType
-> ScalarTypeBase Size Aliasing
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
pn TypeBase Size ()
pt2 PatRetType
ret
typeOf (OpSectionRight QualName VName
_ Info PatType
_ ExpBase Info VName
_ (Info (PName
pn, TypeBase Size ()
pt1), Info (PName, TypeBase Size (), Maybe VName)
_) (Info PatRetType
ret) SrcLoc
_) =
  ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Aliasing
-> PName
-> TypeBase Size ()
-> PatRetType
-> ScalarTypeBase Size Aliasing
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
pn TypeBase Size ()
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 = (TypeBase dim pas -> RetTypeBase dim as -> RetTypeBase dim as)
-> RetTypeBase dim as -> [TypeBase dim pas] -> RetTypeBase dim as
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TypeBase dim pas -> RetTypeBase dim as -> RetTypeBase dim as
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 =
      [VName] -> TypeBase dim as -> RetTypeBase dim as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim as -> RetTypeBase dim as)
-> TypeBase dim as -> RetTypeBase dim as
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty PName
Unnamed (TypeBase dim as -> TypeBase dim ()
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) = TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType TypeBase dim as
t2
   in (TypeBase dim ()
t1 TypeBase dim () -> [TypeBase dim ()] -> [TypeBase dim ()]
forall a. a -> [a] -> [a]
: [TypeBase dim ()]
ps, TypeBase dim ()
r)
unfoldFunType TypeBase dim as
t = ([], TypeBase dim as -> TypeBase dim ()
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], TypeBase Size ())
valBindTypeScheme ValBindBase Info VName
vb =
  ( ValBindBase Info VName -> [TypeParamBase VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [TypeParamBase vn]
valBindTypeParams ValBindBase Info VName
vb,
    [PatBase Info VName] -> StructRetType -> TypeBase Size ()
funType (ValBindBase Info VName -> [PatBase Info VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [PatBase f vn]
valBindParams ValBindBase Info VName
vb) (Info StructRetType -> StructRetType
forall a. Info a -> a
unInfo (ValBindBase Info VName -> Info StructRetType
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 =
  ValBindBase Info VName -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBindBase Info VName
vb
    VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: case ValBindBase Info VName -> [PatBase Info VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [PatBase f vn]
valBindParams ValBindBase Info VName
vb of
      [] -> StructRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims (Info StructRetType -> StructRetType
forall a. Info a -> a
unInfo (ValBindBase Info VName -> Info StructRetType
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 -> TypeBase Size ()
funType [PatBase Info VName]
params StructRetType
ret =
  let RetType [VName]
_ TypeBase Size ()
t = (PatBase Info VName -> StructRetType -> StructRetType)
-> StructRetType -> [PatBase Info VName] -> StructRetType
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((PName, TypeBase Size ()) -> StructRetType -> StructRetType
forall {dim}.
(PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow ((PName, TypeBase Size ()) -> StructRetType -> StructRetType)
-> (PatBase Info VName -> (PName, TypeBase Size ()))
-> PatBase Info VName
-> StructRetType
-> StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> (PName, TypeBase Size ())
patternParam) StructRetType
ret [PatBase Info VName]
params
   in TypeBase Size ()
t
  where
    arrow :: (PName, TypeBase dim ())
-> RetTypeBase dim () -> RetTypeBase dim ()
arrow (PName
xp, TypeBase dim ()
xt) RetTypeBase dim ()
yt = [VName] -> TypeBase dim () -> RetTypeBase dim ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim () -> RetTypeBase dim ())
-> TypeBase dim () -> RetTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ()
-> PName
-> TypeBase dim ()
-> RetTypeBase dim ()
-> ScalarTypeBase dim ()
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 {} -> Set VName
forall a. Monoid a => a
mempty
    Scalar (TypeVar as
_ Uniqueness
_ QualName VName
tn [TypeArg dim]
targs) ->
      [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ VName -> Set VName
forall a. a -> Set a
S.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
tn) Set VName -> [Set VName] -> [Set VName]
forall a. a -> [a] -> [a]
: (TypeArg dim -> Set VName) -> [TypeArg dim] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg dim -> Set VName
forall {dim}. TypeArg dim -> Set VName
typeArgFree [TypeArg dim]
targs
    Scalar (Arrow as
_ PName
_ TypeBase dim ()
t1 (RetType [VName]
_ TypeBase dim as
t2)) -> TypeBase dim () -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim ()
t1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim as
t2
    Scalar (Record Map Name (TypeBase dim as)
fields) -> (TypeBase dim as -> Set VName)
-> Map Name (TypeBase dim as) -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars Map Name (TypeBase dim as)
fields
    Scalar (Sum Map Name [TypeBase dim as]
cs) -> [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (([TypeBase dim as] -> [Set VName])
-> Map Name [TypeBase dim as] -> [Set VName]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (([TypeBase dim as] -> [Set VName])
 -> Map Name [TypeBase dim as] -> [Set VName])
-> ((TypeBase dim as -> Set VName)
    -> [TypeBase dim as] -> [Set VName])
-> (TypeBase dim as -> Set VName)
-> Map Name [TypeBase dim as]
-> [Set VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase dim as -> Set VName) -> [TypeBase dim as] -> [Set VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) TypeBase dim as -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars Map Name [TypeBase dim as]
cs
    Array as
_ Uniqueness
_ Shape dim
_ ScalarTypeBase dim ()
rt -> TypeBase dim () -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars (TypeBase dim () -> Set VName) -> TypeBase dim () -> Set VName
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase dim ()
rt
  where
    typeArgFree :: TypeArg dim -> Set VName
typeArgFree (TypeArgType TypeBase dim ()
ta SrcLoc
_) = TypeBase dim () -> Set VName
forall as dim. Monoid as => TypeBase dim as -> Set VName
typeVars TypeBase dim ()
ta
    typeArgFree TypeArgDim {} = Set VName
forall a. Monoid a => a
mempty

-- | @orderZero t@ is 'True' if the argument type has order 0, i.e., it is not
-- a function type, does not contain a function type as a subcomponent, and may
-- not be instantiated with a function type.
orderZero :: TypeBase dim as -> Bool
orderZero :: 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)) = (TypeBase dim as -> Bool) -> [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim as -> Bool
forall shape as. TypeBase shape as -> Bool
orderZero ([TypeBase dim as] -> Bool) -> [TypeBase dim as] -> Bool
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> [TypeBase dim as]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase dim as)
fs
orderZero (Scalar TypeVar {}) = Bool
True
orderZero (Scalar Arrow {}) = Bool
False
orderZero (Scalar (Sum Map Name [TypeBase dim as]
cs)) = ([TypeBase dim as] -> Bool) -> Map Name [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((TypeBase dim as -> Bool) -> [TypeBase dim as] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim as -> Bool
forall 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
_ -> (PatBase Info vn -> Bool) -> [PatBase Info vn] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all PatBase Info vn -> Bool
forall vn. PatBase Info vn -> Bool
patternOrderZero [PatBase Info vn]
ps
  RecordPat [(Name, PatBase Info vn)]
fs SrcLoc
_ -> ((Name, PatBase Info vn) -> Bool)
-> [(Name, PatBase Info vn)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (PatBase Info vn -> Bool
forall vn. PatBase Info vn -> Bool
patternOrderZero (PatBase Info vn -> Bool)
-> ((Name, PatBase Info vn) -> PatBase Info vn)
-> (Name, PatBase Info vn)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatBase Info vn) -> PatBase Info vn
forall a b. (a, b) -> b
snd) [(Name, PatBase Info vn)]
fs
  PatParens PatBase Info vn
p SrcLoc
_ -> PatBase Info vn -> Bool
forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info vn
p
  Id vn
_ (Info PatType
t) SrcLoc
_ -> PatType -> Bool
forall shape as. TypeBase shape as -> Bool
orderZero PatType
t
  Wildcard (Info PatType
t) SrcLoc
_ -> PatType -> Bool
forall shape as. TypeBase shape as -> Bool
orderZero PatType
t
  PatAscription PatBase Info vn
p TypeExp vn
_ SrcLoc
_ -> PatBase Info vn -> Bool
forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info vn
p
  PatLit PatLit
_ (Info PatType
t) SrcLoc
_ -> PatType -> Bool
forall shape as. TypeBase shape as -> Bool
orderZero PatType
t
  PatConstr Name
_ Info PatType
_ [PatBase Info vn]
ps SrcLoc
_ -> (PatBase Info vn -> Bool) -> [PatBase Info vn] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all PatBase Info vn -> Bool
forall vn. PatBase Info vn -> Bool
patternOrderZero [PatBase Info vn]
ps
  PatAttr AttrInfo vn
_ PatBase Info vn
p SrcLoc
_ -> PatBase Info vn -> Bool
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) = IdentBase f vn -> Set (IdentBase f vn)
forall a. a -> Set a
S.singleton (IdentBase f vn -> Set (IdentBase f vn))
-> IdentBase f vn -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ vn -> f PatType -> SrcLoc -> IdentBase f vn
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
_) = PatBase f vn -> Set (IdentBase f vn)
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
_) = [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase f vn)] -> Set (IdentBase f vn))
-> [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ (PatBase f vn -> Set (IdentBase f vn))
-> [PatBase f vn] -> [Set (IdentBase f vn)]
forall a b. (a -> b) -> [a] -> [b]
map PatBase f vn -> Set (IdentBase f vn)
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
_) = [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase f vn)] -> Set (IdentBase f vn))
-> [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ ((Name, PatBase f vn) -> Set (IdentBase f vn))
-> [(Name, PatBase f vn)] -> [Set (IdentBase f vn)]
forall a b. (a -> b) -> [a] -> [b]
map (PatBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents (PatBase f vn -> Set (IdentBase f vn))
-> ((Name, PatBase f vn) -> PatBase f vn)
-> (Name, PatBase f vn)
-> Set (IdentBase f vn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatBase f vn) -> PatBase f vn
forall a b. (a, b) -> b
snd) [(Name, PatBase f vn)]
fs
patIdents Wildcard {} = Set (IdentBase f vn)
forall a. Monoid a => a
mempty
patIdents (PatAscription PatBase f vn
p TypeExp vn
_ SrcLoc
_) = PatBase f vn -> Set (IdentBase f vn)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents PatBase f vn
p
patIdents PatLit {} = Set (IdentBase f vn)
forall a. Monoid a => a
mempty
patIdents (PatConstr Name
_ f PatType
_ [PatBase f vn]
ps SrcLoc
_) = [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a. Monoid a => [a] -> a
mconcat ([Set (IdentBase f vn)] -> Set (IdentBase f vn))
-> [Set (IdentBase f vn)] -> Set (IdentBase f vn)
forall a b. (a -> b) -> a -> b
$ (PatBase f vn -> Set (IdentBase f vn))
-> [PatBase f vn] -> [Set (IdentBase f vn)]
forall a b. (a -> b) -> [a] -> [b]
map PatBase f vn -> Set (IdentBase f vn)
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
_) = PatBase f vn -> Set (IdentBase f vn)
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
_) = vn -> Set vn
forall a. a -> Set a
S.singleton vn
v
patNames (PatParens PatBase f vn
p SrcLoc
_) = PatBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase f vn
p
patNames (TuplePat [PatBase f vn]
pats SrcLoc
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ (PatBase f vn -> Set vn) -> [PatBase f vn] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map PatBase f vn -> Set vn
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
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ ((Name, PatBase f vn) -> Set vn)
-> [(Name, PatBase f vn)] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map (PatBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames (PatBase f vn -> Set vn)
-> ((Name, PatBase f vn) -> PatBase f vn)
-> (Name, PatBase f vn)
-> Set vn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatBase f vn) -> PatBase f vn
forall a b. (a, b) -> b
snd) [(Name, PatBase f vn)]
fs
patNames Wildcard {} = Set vn
forall a. Monoid a => a
mempty
patNames (PatAscription PatBase f vn
p TypeExp vn
_ SrcLoc
_) = PatBase f vn -> Set vn
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase f vn
p
patNames PatLit {} = Set vn
forall a. Monoid a => a
mempty
patNames (PatConstr Name
_ f PatType
_ [PatBase f vn]
ps SrcLoc
_) = [Set vn] -> Set vn
forall a. Monoid a => [a] -> a
mconcat ([Set vn] -> Set vn) -> [Set vn] -> Set vn
forall a b. (a -> b) -> a -> b
$ (PatBase f vn -> Set vn) -> [PatBase f vn] -> [Set vn]
forall a b. (a -> b) -> [a] -> [b]
map PatBase f vn -> Set vn
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
_) = PatBase f vn -> Set vn
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 =
  [(VName, IdentBase f VName)] -> Map VName (IdentBase f VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, IdentBase f VName)] -> Map VName (IdentBase f VName))
-> [(VName, IdentBase f VName)] -> Map VName (IdentBase f VName)
forall a b. (a -> b) -> a -> b
$ [VName] -> [IdentBase f VName] -> [(VName, IdentBase f VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((IdentBase f VName -> VName) -> [IdentBase f VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map IdentBase f VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName [IdentBase f VName]
idents) [IdentBase f VName]
idents
  where
    idents :: [IdentBase f VName]
idents = Set (IdentBase f VName) -> [IdentBase f VName]
forall a. Set a -> [a]
S.toList (Set (IdentBase f VName) -> [IdentBase f VName])
-> Set (IdentBase f VName) -> [IdentBase f VName]
forall a b. (a -> b) -> a -> b
$ PatBase f VName -> Set (IdentBase f VName)
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
_) = ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ [PatType] -> ScalarTypeBase Size Aliasing
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord ([PatType] -> ScalarTypeBase Size Aliasing)
-> [PatType] -> ScalarTypeBase Size Aliasing
forall a b. (a -> b) -> a -> b
$ (PatBase Info VName -> PatType)
-> [PatBase Info VName] -> [PatType]
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
_) = ScalarTypeBase Size Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size Aliasing -> PatType)
-> ScalarTypeBase Size Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Map Name PatType -> ScalarTypeBase Size Aliasing
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name PatType -> ScalarTypeBase Size Aliasing)
-> Map Name PatType -> ScalarTypeBase Size Aliasing
forall a b. (a -> b) -> a -> b
$ PatBase Info VName -> PatType
patternType (PatBase Info VName -> PatType)
-> Map Name (PatBase Info VName) -> Map Name PatType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Name, PatBase Info VName)] -> Map Name (PatBase Info VName)
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 -> TypeBase Size ()
patternStructType = PatType -> TypeBase Size ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatType -> TypeBase Size ())
-> (PatBase Info VName -> PatType)
-> PatBase Info VName
-> TypeBase Size ()
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, TypeBase Size ())
patternParam (PatParens PatBase Info VName
p SrcLoc
_) =
  PatBase Info VName -> (PName, TypeBase Size ())
patternParam PatBase Info VName
p
patternParam (PatAttr AttrInfo VName
_ PatBase Info VName
p SrcLoc
_) =
  PatBase Info VName -> (PName, TypeBase Size ())
patternParam PatBase Info VName
p
patternParam (PatAscription (Id VName
v (Info PatType
t) SrcLoc
_) TypeExp VName
_ SrcLoc
_) =
  (VName -> PName
Named VName
v, PatType -> TypeBase Size ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)
patternParam (Id VName
v (Info PatType
t) SrcLoc
_) =
  (VName -> PName
Named VName
v, PatType -> TypeBase Size ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)
patternParam PatBase Info VName
p =
  (PName
Unnamed, PatBase Info VName -> TypeBase Size ()
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 =
  [(Name, PrimType)] -> Map Name PrimType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    [ ([Char] -> Name
nameFromString ([Char] -> Name) -> [Char] -> Name
forall a b. (a -> b) -> a -> b
$ PrimType -> [Char]
forall a. Pretty a => a -> [Char]
pretty PrimType
t, PrimType
t)
      | PrimType
t <-
          PrimType
Bool
            PrimType -> [PrimType] -> [PrimType]
forall a. a -> [a] -> [a]
: (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Signed [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
            [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (IntType -> PrimType) -> [IntType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> PrimType
Unsigned [IntType
forall a. Bounded a => a
minBound .. IntType
forall a. Bounded a => a
maxBound]
            [PrimType] -> [PrimType] -> [PrimType]
forall a. [a] -> [a] -> [a]
++ (FloatType -> PrimType) -> [FloatType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map FloatType -> PrimType
FloatType [FloatType
forall a. Bounded a => a
minBound .. FloatType
forall a. Bounded a => a
maxBound]
    ]

-- | The nature of something predefined.  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] -> TypeBase Size () -> Intrinsic
IntrinsicType Liftedness
SizeLifted [Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
t_v SrcLoc
forall a. Monoid a => a
mempty] (TypeBase Size () -> Intrinsic) -> TypeBase Size () -> Intrinsic
forall a b. (a -> b) -> a -> b
$
      ScalarTypeBase Size () -> TypeBase Size ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size () -> TypeBase Size ())
-> ScalarTypeBase Size () -> TypeBase Size ()
forall a b. (a -> b) -> a -> b
$
        ()
-> Uniqueness
-> QualName VName
-> [TypeArg Size]
-> ScalarTypeBase Size ()
forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> QualName VName
forall v. v -> QualName v
qualName VName
acc_v) [TypeArg Size
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 = TypeBase dim () -> SrcLoc -> TypeArg dim
forall dim. TypeBase dim () -> SrcLoc -> TypeArg dim
TypeArgType (ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (()
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim ()
forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> QualName VName
forall v. v -> QualName v
qualName VName
t_v) [])) SrcLoc
forall a. Monoid a => a
mempty

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

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

    t_b :: ScalarTypeBase dim ()
t_b = ()
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim ()
forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar () Uniqueness
Nonunique (VName -> QualName VName
forall v. v -> QualName v
qualName VName
b) []
    arr_b :: Shape dim -> TypeBase dim ()
arr_b Shape dim
s = ()
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique Shape dim
s ScalarTypeBase dim ()
forall {dim}. ScalarTypeBase dim ()
t_b
    uarr_b :: Shape dim -> TypeBase dim ()
uarr_b Shape dim
s = ()
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique Shape dim
s ScalarTypeBase dim ()
forall {dim}. ScalarTypeBase dim ()
t_b
    tp_b :: TypeParamBase VName
tp_b = Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
b SrcLoc
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] = (VName -> TypeParamBase VName) -> [VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) [VName
n, VName
m, VName
k, VName
l, VName
p, VName
q]

    shape :: [VName] -> Shape Size
shape = [Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape ([Size] -> Shape Size)
-> ([VName] -> [Size]) -> [VName] -> Shape Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Size) -> [VName] -> [Size]
forall a b. (a -> b) -> [a] -> [b]
map (QualName VName -> Size
NamedSize (QualName VName -> Size)
-> (VName -> QualName VName) -> VName -> Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
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 =
      ()
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array
        ()
        Uniqueness
Nonunique
        Shape dim
s
        (Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record ([(Name, TypeBase dim ())] -> Map Name (TypeBase dim ())
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase dim ())] -> Map Name (TypeBase dim ()))
-> [(Name, TypeBase dim ())] -> Map Name (TypeBase dim ())
forall a b. (a -> b) -> a -> b
$ [Name] -> [TypeBase dim ()] -> [(Name, TypeBase dim ())]
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 = TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
forall {dim}.
TypeBase dim () -> TypeBase dim () -> Shape dim -> TypeBase dim ()
tuple_arr TypeBase dim ()
x TypeBase dim ()
y Shape dim
s TypeBase dim () -> Uniqueness -> TypeBase dim ()
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 = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty PName
Unnamed TypeBase dim ()
x ([VName] -> TypeBase dim as -> RetTypeBase dim as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase dim as
y)

    arr_ka :: TypeBase Size ()
arr_ka = ()
-> Uniqueness
-> Shape Size
-> ScalarTypeBase Size ()
-> TypeBase Size ()
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [QualName VName -> Size
NamedSize (QualName VName -> Size) -> QualName VName -> Size
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
k]) ScalarTypeBase Size ()
forall {dim}. ScalarTypeBase dim ()
t_a
    uarr_ka :: TypeBase Size ()
uarr_ka = ()
-> Uniqueness
-> Shape Size
-> ScalarTypeBase Size ()
-> TypeBase Size ()
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Unique ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [QualName VName -> Size
NamedSize (QualName VName -> Size) -> QualName VName -> Size
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
k]) ScalarTypeBase Size ()
forall {dim}. ScalarTypeBase dim ()
t_a
    arr_kb :: TypeBase Size ()
arr_kb = ()
-> Uniqueness
-> Shape Size
-> ScalarTypeBase Size ()
-> TypeBase Size ()
forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
Nonunique ([Size] -> Shape Size
forall dim. [dim] -> Shape dim
Shape [QualName VName -> Size
NamedSize (QualName VName -> Size) -> QualName VName -> Size
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
k]) ScalarTypeBase Size ()
forall {dim}. ScalarTypeBase dim ()
t_b
    karr :: TypeBase dim () -> TypeBase dim as -> TypeBase dim as
karr TypeBase dim ()
x TypeBase dim as
y = ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
forall a. Monoid a => a
mempty (VName -> PName
Named VName
k) TypeBase dim ()
x ([VName] -> TypeBase dim as -> RetTypeBase dim as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase dim as
y)

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

    namify :: Int -> ([Char], b) -> (VName, b)
namify Int
i ([Char]
x, b
y) = (Name -> Int -> VName
VName ([Char] -> Name
nameFromString [Char]
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 ((PrimType -> PrimType) -> [PrimType] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> PrimType
unPrim [PrimType]
ts) (PrimType -> Intrinsic) -> PrimType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
unPrim PrimType
t)

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

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

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

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

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

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

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

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

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

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

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

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

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

-- | Is this file part of the built-in prelude?
isBuiltin :: FilePath -> Bool
isBuiltin :: [Char] -> Bool
isBuiltin = ([Char]
"/prelude/" [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf`)

-- | 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 a -> Loc
forall a. Located a => a -> Loc
locOf a
x of
    Loc
NoLoc -> Bool
False
    Loc Pos
pos Pos
_ -> [Char] -> Bool
isBuiltin ([Char] -> Bool) -> [Char] -> Bool
forall a b. (a -> b) -> a -> b
$ Pos -> [Char]
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 = [Int] -> Int
forall a (f :: * -> *). (Num a, Ord a, Foldable f) => f a -> a
maxinum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (VName -> Int) -> [VName] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Int
baseTag ([VName] -> [Int]) -> [VName] -> [Int]
forall a b. (a -> b) -> a -> b
$ Map VName Intrinsic -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName Intrinsic
intrinsics

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

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

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

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

modExpImports :: ModExpBase f vn -> [(String, SrcLoc)]
modExpImports :: forall (f :: * -> *) vn. ModExpBase f vn -> [([Char], SrcLoc)]
modExpImports ModVar {} = []
modExpImports (ModParens ModExpBase f vn
p SrcLoc
_) = ModExpBase f vn -> [([Char], SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [([Char], SrcLoc)]
modExpImports ModExpBase f vn
p
modExpImports (ModImport [Char]
f f [Char]
_ SrcLoc
loc) = [([Char]
f, SrcLoc
loc)]
modExpImports (ModDecs [DecBase f vn]
ds SrcLoc
_) = (DecBase f vn -> [([Char], SrcLoc)])
-> [DecBase f vn] -> [([Char], SrcLoc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [([Char], SrcLoc)]
forall (f :: * -> *) vn. DecBase f vn -> [([Char], SrcLoc)]
decImports [DecBase f vn]
ds
modExpImports (ModApply ModExpBase f vn
_ ModExpBase f vn
me f (Map VName VName)
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [([Char], SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [([Char], SrcLoc)]
modExpImports ModExpBase f vn
me
modExpImports (ModAscript ModExpBase f vn
me SigExpBase f vn
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [([Char], SrcLoc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [([Char], 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 = (VName -> Set VName) -> Set VName -> Set VName
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 = VName -> Set VName
forall a. a -> Set a
S.singleton VName
v Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
-> (Set VName -> Set VName) -> Maybe (Set VName) -> Set VName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set VName
forall a. Monoid a => a
mempty ((VName -> Set VName) -> Set VName -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> Set VName
reach) (VName -> Map VName (Set VName) -> Maybe (Set VName)
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 = (DecBase Info VName -> Map VName (Set VName))
-> [DecBase Info VName] -> Map VName (Set VName)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase Info VName -> Map VName (Set VName)
forall {a} {f :: * -> *}. Ord a => DecBase f a -> Map a (Set a)
onDec ([DecBase Info VName] -> Map VName (Set VName))
-> [DecBase Info VName] -> Map VName (Set VName)
forall a b. (a -> b) -> a -> b
$ ProgBase Info VName -> [DecBase Info VName]
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 {} = Map a (Set a)
forall a. Monoid a => a
mempty
        onDec ModDec {} = Map a (Set a)
forall a. Monoid a => a
mempty
        onDec (SigDec SigBindBase f a
sb) =
          a -> Set a -> Map a (Set a)
forall k a. k -> a -> Map k a
M.singleton (SigBindBase f a -> a
forall (f :: * -> *) vn. SigBindBase f vn -> vn
sigName SigBindBase f a
sb) (SigExpBase f a -> Set a
forall {a} {f :: * -> *}. Ord a => SigExpBase f a -> Set a
onSigExp (SigBindBase f a -> SigExpBase f a
forall (f :: * -> *) vn. SigBindBase f vn -> SigExpBase f vn
sigExp SigBindBase f a
sb))
        onDec TypeDec {} = Map a (Set a)
forall a. Monoid a => a
mempty
        onDec ValDec {} = Map a (Set a)
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 {} = Map a (Set a)
forall a. Monoid a => a
mempty

        onSigExp :: SigExpBase f a -> Set a
onSigExp (SigVar QualName a
v f (Map VName VName)
_ SrcLoc
_) = a -> Set a
forall a. a -> Set a
S.singleton (a -> Set a) -> a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
v
        onSigExp (SigParens SigExpBase f a
e SrcLoc
_) = SigExpBase f a -> Set a
onSigExp SigExpBase f a
e
        onSigExp (SigSpecs [SpecBase f a]
ss SrcLoc
_) = (SpecBase f a -> Set a) -> [SpecBase f a] -> Set a
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 Set a -> Set a -> Set a
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 {} = Set a
forall a. Monoid a => a
mempty
        onSpec TypeSpec {} = Set a
forall a. Monoid a => a
mempty
        onSpec TypeAbbrSpec {} = Set a
forall a. Monoid a => a
mempty
        onSpec (ModSpec a
vn SigExpBase f a
e Maybe DocComment
_ SrcLoc
_) = a -> Set a
forall a. a -> Set a
S.singleton a
vn Set a -> Set a -> Set a
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 = (DecBase Info VName -> Set VName)
-> [DecBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase Info VName -> Set VName
forall {f :: * -> *}. DecBase f VName -> Set VName
onDec ([DecBase Info VName] -> Set VName)
-> [DecBase Info VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ ProgBase Info VName -> [DecBase Info VName]
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) =
          Set VName
-> ((SigExpBase f VName, f (Map VName VName)) -> Set VName)
-> Maybe (SigExpBase f VName, f (Map VName VName))
-> Set VName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set VName
forall a. Monoid a => a
mempty (SigExpBase f VName -> Set VName
forall {a} {f :: * -> *}. Ord a => SigExpBase f a -> Set a
onSigExp (SigExpBase f VName -> Set VName)
-> ((SigExpBase f VName, f (Map VName VName))
    -> SigExpBase f VName)
-> (SigExpBase f VName, f (Map VName VName))
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SigExpBase f VName, f (Map VName VName)) -> SigExpBase f VName
forall a b. (a, b) -> a
fst) (ModBindBase f VName
-> Maybe (SigExpBase f VName, f (Map VName VName))
forall (f :: * -> *) vn.
ModBindBase f vn -> Maybe (SigExpBase f vn, f (Map VName VName))
modSignature ModBindBase f VName
md) Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> ModExpBase f VName -> Set VName
onModExp (ModBindBase f VName -> ModExpBase f VName
forall (f :: * -> *) vn. ModBindBase f vn -> ModExpBase f vn
modExp ModBindBase f VName
md)
        onDec SigDec {} = Set VName
forall a. Monoid a => a
mempty
        onDec TypeDec {} = Set VName
forall a. Monoid a => a
mempty
        onDec ValDec {} = Set VName
forall a. Monoid a => a
mempty
        onDec LocalDec {} = Set VName
forall a. Monoid a => a
mempty
        onDec ImportDec {} = Set VName
forall a. Monoid a => a
mempty

        onModExp :: ModExpBase f VName -> Set VName
onModExp ModVar {} = Set VName
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 {} = Set VName
forall a. Monoid a => a
mempty
        onModExp (ModDecs [DecBase f VName]
ds SrcLoc
_) = [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (DecBase f VName -> Set VName) -> [DecBase f VName] -> [Set VName]
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 Set VName -> Set VName -> Set VName
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 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> SigExpBase f VName -> Set VName
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
_) =
          ModParamBase f VName -> Set VName
forall {f :: * -> *}. ModParamBase f VName -> Set VName
onModParam ModParamBase f VName
p Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
-> ((SigExpBase f VName, f (Map VName VName)) -> Set VName)
-> Maybe (SigExpBase f VName, f (Map VName VName))
-> Set VName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set VName
forall a. Monoid a => a
mempty (SigExpBase f VName -> Set VName
forall {a} {f :: * -> *}. Ord a => SigExpBase f a -> Set a
onSigExp (SigExpBase f VName -> Set VName)
-> ((SigExpBase f VName, f (Map VName VName))
    -> SigExpBase f VName)
-> (SigExpBase f VName, f (Map VName VName))
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SigExpBase f VName, f (Map VName VName)) -> SigExpBase f VName
forall a b. (a, b) -> a
fst) Maybe (SigExpBase f VName, f (Map VName VName))
r Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> ModExpBase f VName -> Set VName
onModExp ModExpBase f VName
me

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

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

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

-- | Given an operator name, return the operator that determines its
-- syntactical properties.
leadingOperator :: Name -> BinOp
leadingOperator :: Name -> BinOp
leadingOperator Name
s =
  BinOp
-> (([Char], BinOp) -> BinOp) -> Maybe ([Char], BinOp) -> BinOp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe BinOp
Backtick ([Char], BinOp) -> BinOp
forall a b. (a, b) -> b
snd (Maybe ([Char], BinOp) -> BinOp) -> Maybe ([Char], BinOp) -> BinOp
forall a b. (a -> b) -> a -> b
$
    (([Char], BinOp) -> Bool)
-> [([Char], BinOp)] -> Maybe ([Char], BinOp)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
s') ([Char] -> Bool)
-> (([Char], BinOp) -> [Char]) -> ([Char], BinOp) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char], BinOp) -> [Char]
forall a b. (a, b) -> a
fst) ([([Char], BinOp)] -> Maybe ([Char], BinOp))
-> [([Char], BinOp)] -> Maybe ([Char], BinOp)
forall a b. (a -> b) -> a -> b
$
      (([Char], BinOp) -> Down Int)
-> [([Char], BinOp)] -> [([Char], BinOp)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int -> Down Int
forall a. a -> Down a
Down (Int -> Down Int)
-> (([Char], BinOp) -> Int) -> ([Char], BinOp) -> Down Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Char] -> Int)
-> (([Char], BinOp) -> [Char]) -> ([Char], BinOp) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char], BinOp) -> [Char]
forall a b. (a, b) -> a
fst) ([([Char], BinOp)] -> [([Char], BinOp)])
-> [([Char], BinOp)] -> [([Char], BinOp)]
forall a b. (a -> b) -> a -> b
$
        [[Char]] -> [BinOp] -> [([Char], BinOp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((BinOp -> [Char]) -> [BinOp] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map BinOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty [BinOp]
operators) [BinOp]
operators
  where
    s' :: [Char]
s' = Name -> [Char]
nameToString Name
s
    operators :: [BinOp]
    operators :: [BinOp]
operators = [BinOp
forall a. Bounded a => a
minBound .. BinOp
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, TypeBase Size ())]
progHoles = (DecBase Info VName -> [(Loc, TypeBase Size ())])
-> [DecBase Info VName] -> [(Loc, TypeBase Size ())]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase Info VName -> [(Loc, TypeBase Size ())]
holesInDec ([DecBase Info VName] -> [(Loc, TypeBase Size ())])
-> (ProgBase Info VName -> [DecBase Info VName])
-> ProgBase Info VName
-> [(Loc, TypeBase Size ())]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgBase Info VName -> [DecBase Info VName]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs
  where
    holesInDec :: DecBase Info VName -> [(Loc, TypeBase Size ())]
holesInDec (ValDec ValBindBase Info VName
vb) = ExpBase Info VName -> [(Loc, TypeBase Size ())]
holesInExp (ExpBase Info VName -> [(Loc, TypeBase Size ())])
-> ExpBase Info VName -> [(Loc, TypeBase Size ())]
forall a b. (a -> b) -> a -> b
$ ValBindBase Info VName -> ExpBase Info VName
forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBindBase Info VName
vb
    holesInDec (ModDec ModBindBase Info VName
me) = ModExpBase Info VName -> [(Loc, TypeBase Size ())]
holesInModExp (ModExpBase Info VName -> [(Loc, TypeBase Size ())])
-> ModExpBase Info VName -> [(Loc, TypeBase Size ())]
forall a b. (a -> b) -> a -> b
$ ModBindBase Info VName -> ModExpBase Info VName
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, TypeBase Size ())]
holesInModExp ModExpBase Info VName
me
    holesInDec (LocalDec DecBase Info VName
d SrcLoc
_) = DecBase Info VName -> [(Loc, TypeBase Size ())]
holesInDec DecBase Info VName
d
    holesInDec TypeDec {} = [(Loc, TypeBase Size ())]
forall a. Monoid a => a
mempty
    holesInDec SigDec {} = [(Loc, TypeBase Size ())]
forall a. Monoid a => a
mempty
    holesInDec ImportDec {} = [(Loc, TypeBase Size ())]
forall a. Monoid a => a
mempty

    holesInModExp :: ModExpBase Info VName -> [(Loc, TypeBase Size ())]
holesInModExp (ModDecs [DecBase Info VName]
ds SrcLoc
_) = (DecBase Info VName -> [(Loc, TypeBase Size ())])
-> [DecBase Info VName] -> [(Loc, TypeBase Size ())]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase Info VName -> [(Loc, TypeBase Size ())]
holesInDec [DecBase Info VName]
ds
    holesInModExp (ModParens ModExpBase Info VName
me SrcLoc
_) = ModExpBase Info VName -> [(Loc, TypeBase Size ())]
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, TypeBase Size ())]
holesInModExp ModExpBase Info VName
x [(Loc, TypeBase Size ())]
-> [(Loc, TypeBase Size ())] -> [(Loc, TypeBase Size ())]
forall a. Semigroup a => a -> a -> a
<> ModExpBase Info VName -> [(Loc, TypeBase Size ())]
holesInModExp ModExpBase Info VName
y
    holesInModExp (ModAscript ModExpBase Info VName
me SigExpBase Info VName
_ Info (Map VName VName)
_ SrcLoc
_) = ModExpBase Info VName -> [(Loc, TypeBase Size ())]
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, TypeBase Size ())]
holesInModExp ModExpBase Info VName
me
    holesInModExp ModVar {} = [(Loc, TypeBase Size ())]
forall a. Monoid a => a
mempty
    holesInModExp ModImport {} = [(Loc, TypeBase Size ())]
forall a. Monoid a => a
mempty

    holesInExp :: ExpBase Info VName -> [(Loc, TypeBase Size ())]
holesInExp = (State [(Loc, TypeBase Size ())] (ExpBase Info VName)
 -> [(Loc, TypeBase Size ())] -> [(Loc, TypeBase Size ())])
-> [(Loc, TypeBase Size ())]
-> State [(Loc, TypeBase Size ())] (ExpBase Info VName)
-> [(Loc, TypeBase Size ())]
forall a b c. (a -> b -> c) -> b -> a -> c
flip State [(Loc, TypeBase Size ())] (ExpBase Info VName)
-> [(Loc, TypeBase Size ())] -> [(Loc, TypeBase Size ())]
forall s a. State s a -> s -> s
execState [(Loc, TypeBase Size ())]
forall a. Monoid a => a
mempty (State [(Loc, TypeBase Size ())] (ExpBase Info VName)
 -> [(Loc, TypeBase Size ())])
-> (ExpBase Info VName
    -> State [(Loc, TypeBase Size ())] (ExpBase Info VName))
-> ExpBase Info VName
-> [(Loc, TypeBase Size ())]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExpBase Info VName
-> State [(Loc, TypeBase Size ())] (ExpBase Info VName)
forall {m :: * -> *}.
MonadState [(Loc, TypeBase Size ())] 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
      ([(Loc, TypeBase Size ())] -> [(Loc, TypeBase Size ())]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc, PatType -> TypeBase Size ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) (Loc, TypeBase Size ())
-> [(Loc, TypeBase Size ())] -> [(Loc, TypeBase Size ())]
forall a. a -> [a] -> [a]
:)
      ExpBase Info VName -> m (ExpBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase Info VName
e
    onExp ExpBase Info VName
e = ASTMapper m -> ExpBase Info VName -> m (ExpBase Info VName)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (ASTMapper m
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