-- | 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,
    intrinsicVar,
    maxIntrinsicTag,
    isIntrinsic,
    namesToPrimTypes,
    qualName,
    qualify,
    primValueType,
    leadingOperator,
    progImports,
    decImports,
    progModuleTypes,
    identifierReference,
    prettyStacktrace,
    progHoles,
    defaultEntryPoint,
    paramName,
    anySize,
    isAnySize,
    setApplyLoc,

    -- * Queries on expressions
    typeOf,
    valBindTypeScheme,
    valBindBound,
    funType,
    stripExp,
    subExps,
    similarExps,
    sameExp,

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

    -- * Queries on types
    freshness,
    fresh,
    diet,
    arrayRank,
    arrayShape,
    orderZero,
    unfoldFunType,
    foldFunType,
    typeQualVars,
    typeVars,
    isAccType,

    -- * Operations on types
    peelArray,
    stripArray,
    arrayOf,
    arrayOfWithAliases,
    toStructural,
    toStruct,
    toRes,
    toParam,
    resToParam,
    paramToRes,
    toResRet,
    setMode,
    noSizes,
    traverseDims,
    DimPos (..),
    tupleRecord,
    isTupleRecord,
    areTupleFields,
    tupleFields,
    tupleFieldNames,
    sortFields,
    sortConstrs,
    isTypeParam,
    isSizeParam,
    matchDims,

    -- * Un-typechecked ASTs
    UncheckedType,
    UncheckedTypeExp,
    UncheckedIdent,
    UncheckedDimIndex,
    UncheckedSlice,
    UncheckedExp,
    UncheckedModExp,
    UncheckedModTypeExp,
    UncheckedTypeParam,
    UncheckedPat,
    UncheckedValBind,
    UncheckedTypeBind,
    UncheckedModTypeBind,
    UncheckedModBind,
    UncheckedDec,
    UncheckedSpec,
    UncheckedProg,
    UncheckedCase,

    -- * Type-checked ASTs
    Ident,
    DimIndex,
    Slice,
    AppExp,
    Exp,
    Pat,
    ModExp,
    ModParam,
    ModTypeExp,
    ModBind,
    ModTypeBind,
    ValBind,
    Dec,
    Spec,
    Prog,
    TypeBind,
    StructTypeArg,
    ScalarType,
    TypeParam,
    Case,
  )
where

import Control.Monad
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable (bitraverse)
import Data.Char
import Data.Foldable
import Data.List (genericLength, isPrefixOf, sortOn)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.Util (maxinum)
import Futhark.Util.Pretty
import Language.Futhark.Primitive qualified as Primitive
import Language.Futhark.Syntax
import Language.Futhark.Traversals
import Language.Futhark.Tuple

-- | The name of the default program entry point (@main@).
defaultEntryPoint :: Name
defaultEntryPoint :: Name
defaultEntryPoint = String -> Name
nameFromString String
"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 d o -> Int
arrayRank :: forall d o. TypeBase d o -> Int
arrayRank = Shape d -> Int
forall dim. Shape dim -> Int
shapeRank (Shape d -> Int)
-> (TypeBase d o -> Shape d) -> TypeBase d o -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase d o -> Shape d
forall dim o. TypeBase dim o -> Shape dim
arrayShape

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

-- | Change the shape of a type to be just the rank.
noSizes :: TypeBase Size o -> TypeBase () o
noSizes :: forall o. TypeBase Exp o -> TypeBase () o
noSizes = (Exp -> ()) -> TypeBase Exp o -> TypeBase () o
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Exp -> ()) -> TypeBase Exp o -> TypeBase () o)
-> (Exp -> ()) -> TypeBase Exp o -> TypeBase () o
forall a b. (a -> b) -> a -> b
$ () -> Exp -> ()
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
$c== :: DimPos -> DimPos -> Bool
== :: DimPos -> DimPos -> Bool
$c/= :: DimPos -> DimPos -> Bool
/= :: 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
$ccompare :: DimPos -> DimPos -> Ordering
compare :: DimPos -> DimPos -> Ordering
$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
>= :: DimPos -> DimPos -> Bool
$cmax :: DimPos -> DimPos -> DimPos
max :: DimPos -> DimPos -> DimPos
$cmin :: DimPos -> DimPos -> DimPos
min :: DimPos -> DimPos -> DimPos
Ord, Int -> DimPos -> ShowS
[DimPos] -> ShowS
DimPos -> String
(Int -> DimPos -> ShowS)
-> (DimPos -> String) -> ([DimPos] -> ShowS) -> Show DimPos
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DimPos -> ShowS
showsPrec :: Int -> DimPos -> ShowS
$cshow :: DimPos -> String
show :: DimPos -> String
$cshowList :: [DimPos] -> ShowS
showList :: [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 o.
  (Applicative f) =>
  (S.Set VName -> DimPos -> fdim -> f tdim) ->
  TypeBase fdim o ->
  f (TypeBase tdim o)
traverseDims :: forall (f :: * -> *) fdim tdim o.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim o -> f (TypeBase tdim o)
traverseDims Set VName -> DimPos -> fdim -> f tdim
f = Set VName -> DimPos -> TypeBase fdim o -> f (TypeBase tdim o)
forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
forall a. Monoid a => a
mempty DimPos
PosImmediate
  where
    go ::
      forall o'.
      S.Set VName ->
      DimPos ->
      TypeBase fdim o' ->
      f (TypeBase tdim o')
    go :: forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
bound DimPos
b t :: TypeBase fdim o'
t@Array {} =
      (fdim -> f tdim)
-> (o' -> f o') -> TypeBase fdim o' -> f (TypeBase tdim o')
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
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) o' -> f o'
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase fdim o'
t
    go Set VName
bound DimPos
b (Scalar (Record Map Name (TypeBase fdim o')
fields)) =
      ScalarTypeBase tdim o' -> TypeBase tdim o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase tdim o' -> TypeBase tdim o')
-> (Map Name (TypeBase tdim o') -> ScalarTypeBase tdim o')
-> Map Name (TypeBase tdim o')
-> TypeBase tdim o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase tdim o') -> ScalarTypeBase tdim o'
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase tdim o') -> TypeBase tdim o')
-> f (Map Name (TypeBase tdim o')) -> f (TypeBase tdim o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase fdim o' -> f (TypeBase tdim o'))
-> Map Name (TypeBase fdim o') -> f (Map Name (TypeBase tdim o'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
bound DimPos
b) Map Name (TypeBase fdim o')
fields
    go Set VName
bound DimPos
b (Scalar (TypeVar o'
as QualName VName
tn [TypeArg fdim]
targs)) =
      ScalarTypeBase tdim o' -> TypeBase tdim o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase tdim o' -> TypeBase tdim o')
-> f (ScalarTypeBase tdim o') -> f (TypeBase tdim o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (o' -> QualName VName -> [TypeArg tdim] -> ScalarTypeBase tdim o'
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o'
as QualName VName
tn ([TypeArg tdim] -> ScalarTypeBase tdim o')
-> f [TypeArg tdim] -> f (ScalarTypeBase tdim o')
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)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (QualName VName
-> Set VName -> DimPos -> TypeArg fdim -> f (TypeArg tdim)
onTypeArg QualName VName
tn Set VName
bound DimPos
b) [TypeArg fdim]
targs)
    go Set VName
bound DimPos
b (Scalar (Sum Map Name [TypeBase fdim o']
cs)) =
      ScalarTypeBase tdim o' -> TypeBase tdim o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase tdim o' -> TypeBase tdim o')
-> (Map Name [TypeBase tdim o'] -> ScalarTypeBase tdim o')
-> Map Name [TypeBase tdim o']
-> TypeBase tdim o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase tdim o'] -> ScalarTypeBase tdim o'
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase tdim o'] -> TypeBase tdim o')
-> f (Map Name [TypeBase tdim o']) -> f (TypeBase tdim o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase fdim o'] -> f [TypeBase tdim o'])
-> Map Name [TypeBase fdim o'] -> f (Map Name [TypeBase tdim o'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse ((TypeBase fdim o' -> f (TypeBase tdim o'))
-> [TypeBase fdim o'] -> f [TypeBase tdim o']
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
bound DimPos
b)) Map Name [TypeBase fdim o']
cs
    go Set VName
_ DimPos
_ (Scalar (Prim PrimType
t)) =
      TypeBase tdim o' -> f (TypeBase tdim o')
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase tdim o' -> f (TypeBase tdim o'))
-> TypeBase tdim o' -> f (TypeBase tdim o')
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase tdim o' -> TypeBase tdim o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase tdim o' -> TypeBase tdim o')
-> ScalarTypeBase tdim o' -> TypeBase tdim o'
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase tdim o'
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
t
    go Set VName
bound DimPos
_ (Scalar (Arrow o'
als PName
p Diet
o TypeBase fdim NoMode
t1 (RetType [VName]
dims TypeBase fdim Freshness
t2))) =
      ScalarTypeBase tdim o' -> TypeBase tdim o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase tdim o' -> TypeBase tdim o')
-> f (ScalarTypeBase tdim o') -> f (TypeBase tdim o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (o'
-> PName
-> Diet
-> TypeBase tdim NoMode
-> RetTypeBase tdim Freshness
-> ScalarTypeBase tdim o'
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o'
als PName
p Diet
o (TypeBase tdim NoMode
 -> RetTypeBase tdim Freshness -> ScalarTypeBase tdim o')
-> f (TypeBase tdim NoMode)
-> f (RetTypeBase tdim Freshness -> ScalarTypeBase tdim o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> DimPos -> TypeBase fdim NoMode -> f (TypeBase tdim NoMode)
forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
bound' DimPos
PosParam TypeBase fdim NoMode
t1 f (RetTypeBase tdim Freshness -> ScalarTypeBase tdim o')
-> f (RetTypeBase tdim Freshness) -> f (ScalarTypeBase tdim o')
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([VName] -> TypeBase tdim Freshness -> RetTypeBase tdim Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims (TypeBase tdim Freshness -> RetTypeBase tdim Freshness)
-> f (TypeBase tdim Freshness) -> f (RetTypeBase tdim Freshness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> DimPos -> TypeBase fdim Freshness -> f (TypeBase tdim Freshness)
forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
bound' DimPos
PosReturn TypeBase fdim Freshness
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 :: QualName VName
-> Set VName -> DimPos -> TypeArg fdim -> f (TypeArg tdim)
onTypeArg QualName VName
_ Set VName
bound DimPos
b (TypeArgDim fdim
d) =
      tdim -> TypeArg tdim
forall dim. dim -> TypeArg dim
TypeArgDim (tdim -> TypeArg tdim) -> f tdim -> f (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
    onTypeArg QualName VName
tn Set VName
bound DimPos
b (TypeArgType TypeBase fdim NoMode
t) =
      TypeBase tdim NoMode -> TypeArg tdim
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (TypeBase tdim NoMode -> TypeArg tdim)
-> f (TypeBase tdim NoMode) -> f (TypeArg tdim)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> DimPos -> TypeBase fdim NoMode -> f (TypeBase tdim NoMode)
forall o'.
Set VName -> DimPos -> TypeBase fdim o' -> f (TypeBase tdim o')
go Set VName
bound DimPos
b' TypeBase fdim NoMode
t
      where
        b' :: DimPos
b' =
          if QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
tn VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== (VName, Intrinsic) -> VName
forall a b. (a, b) -> a
fst (VName, Intrinsic)
intrinsicAcc
            then DimPos
b
            else DimPos
PosParam

-- | Return the freshness of a type.  A composite type is fresh if any
-- of its components are.
freshness :: TypeBase shape Freshness -> Freshness
freshness :: forall shape. TypeBase shape Freshness -> Freshness
freshness (Array Freshness
o Shape shape
_ ScalarTypeBase shape NoMode
_) = Freshness
o
freshness (Scalar (TypeVar Freshness
o QualName VName
_ [TypeArg shape]
_)) = Freshness
o
freshness (Scalar (Sum Map Name [TypeBase shape Freshness]
ts))
  | ([TypeBase shape Freshness] -> Bool)
-> Map Name [TypeBase shape Freshness] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((TypeBase shape Freshness -> Bool)
-> [TypeBase shape Freshness] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeBase shape Freshness -> Bool
forall shape. TypeBase shape Freshness -> Bool
fresh) Map Name [TypeBase shape Freshness]
ts = Freshness
Fresh
freshness (Scalar (Record Map Name (TypeBase shape Freshness)
fs))
  | (TypeBase shape Freshness -> Bool)
-> Map Name (TypeBase shape Freshness) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeBase shape Freshness -> Bool
forall shape. TypeBase shape Freshness -> Bool
fresh Map Name (TypeBase shape Freshness)
fs = Freshness
Fresh
freshness TypeBase shape Freshness
_ = Freshness
Nonfresh

-- | @fresh t@ is 'True' if the type of the argument is fresh.
fresh :: TypeBase shape Freshness -> Bool
fresh :: forall shape. TypeBase shape Freshness -> Bool
fresh = (Freshness -> Freshness -> Bool
forall a. Eq a => a -> a -> Bool
== Freshness
Fresh) (Freshness -> Bool)
-> (TypeBase shape Freshness -> Freshness)
-> TypeBase shape Freshness
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase shape Freshness -> Freshness
forall shape. TypeBase shape Freshness -> Freshness
freshness

-- | @diet t@ returns a description of how a function parameter of
-- type @t@ consumes its argument.
diet :: TypeBase shape Diet -> Diet
diet :: forall shape. TypeBase shape Diet -> Diet
diet (Scalar (Record Map Name (TypeBase shape Diet)
ets)) = (Diet -> Diet -> Diet) -> Diet -> Map Name Diet -> Diet
forall b a. (b -> a -> b) -> b -> Map Name a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Diet -> Diet -> Diet
forall a. Ord a => a -> a -> a
max Diet
Observe (Map Name Diet -> Diet) -> Map Name Diet -> Diet
forall a b. (a -> b) -> a -> b
$ (TypeBase shape Diet -> Diet)
-> Map Name (TypeBase shape Diet) -> Map Name Diet
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet Map Name (TypeBase shape Diet)
ets
diet (Scalar (Prim PrimType
_)) = Diet
Observe
diet (Scalar (Arrow {})) = Diet
Observe
diet (Array Diet
d Shape shape
_ ScalarTypeBase shape NoMode
_) = Diet
d
diet (Scalar (TypeVar Diet
d QualName VName
_ [TypeArg shape]
_)) = Diet
d
diet (Scalar (Sum Map Name [TypeBase shape Diet]
cs)) = (Diet -> Diet -> Diet) -> Diet -> [Diet] -> Diet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Diet -> Diet -> Diet
forall a. Ord a => a -> a -> a
max Diet
Observe ([Diet] -> Diet) -> [Diet] -> Diet
forall a b. (a -> b) -> a -> b
$ ([TypeBase shape Diet] -> [Diet])
-> Map Name [TypeBase shape Diet] -> [Diet]
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((TypeBase shape Diet -> Diet) -> [TypeBase shape Diet] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase shape Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet) Map Name [TypeBase shape Diet]
cs

-- | Convert any type to one that has rank information, no alias
-- information, and no embedded names.
toStructural ::
  TypeBase dim o ->
  TypeBase () NoMode
toStructural :: forall dim o. TypeBase dim o -> TypeBase () NoMode
toStructural = (dim -> ())
-> (o -> NoMode) -> TypeBase dim o -> TypeBase () NoMode
forall a b c d.
(a -> b) -> (c -> d) -> TypeBase a c -> TypeBase b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (() -> dim -> ()
forall a b. a -> b -> a
const ()) (NoMode -> o -> NoMode
forall a b. a -> b -> a
const NoMode
NoMode)

-- | Remove the annotation from a type.
toStruct :: TypeBase dim o -> TypeBase dim NoMode
toStruct :: forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct = (o -> NoMode) -> TypeBase dim o -> TypeBase dim NoMode
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (NoMode -> o -> NoMode
forall a b. a -> b -> a
const NoMode
NoMode)

-- | Uses 'Observe'.
toParam :: Diet -> TypeBase Size o -> ParamType
toParam :: forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
d = (o -> Diet) -> TypeBase Exp o -> ParamType
forall a b. (a -> b) -> TypeBase Exp a -> TypeBase Exp b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Diet -> o -> Diet
forall a b. a -> b -> a
const Diet
d)

-- | Convert to 'ResType'
toRes :: Freshness -> TypeBase Size o -> ResType
toRes :: forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
o = (o -> Freshness) -> TypeBase Exp o -> ResType
forall a b. (a -> b) -> TypeBase Exp a -> TypeBase Exp b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Freshness -> o -> Freshness
forall a b. a -> b -> a
const Freshness
o)

-- | Convert to 'ResRetType'
toResRet :: Freshness -> RetTypeBase Size o -> ResRetType
toResRet :: forall o. Freshness -> RetTypeBase Exp o -> ResRetType
toResRet Freshness
o = (o -> Freshness) -> RetTypeBase Exp o -> ResRetType
forall b c a. (b -> c) -> RetTypeBase a b -> RetTypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Freshness -> o -> Freshness
forall a b. a -> b -> a
const Freshness
o)

-- | Preserves relation between 'Diet' and 'Freshness'.
resToParam :: ResType -> ParamType
resToParam :: ResType -> ParamType
resToParam = (Freshness -> Diet) -> ResType -> ParamType
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Freshness -> Diet
f
  where
    f :: Freshness -> Diet
f Freshness
Fresh = Diet
Consume
    f Freshness
Nonfresh = Diet
Observe

-- | Preserves relation between 'Diet' and 'Freshness'.
paramToRes :: ParamType -> ResType
paramToRes :: ParamType -> ResType
paramToRes = (Diet -> Freshness) -> ParamType -> ResType
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Diet -> Freshness
f
  where
    f :: Diet -> Freshness
f Diet
Consume = Freshness
Fresh
    f Diet
Observe = Freshness
Nonfresh

-- | @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 o -> Maybe (TypeBase dim o)
peelArray :: forall dim o. Int -> TypeBase dim o -> Maybe (TypeBase dim o)
peelArray Int
n (Array o
o Shape dim
shape ScalarTypeBase dim NoMode
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 o -> Maybe (TypeBase dim o)
forall a. a -> Maybe a
Just (TypeBase dim o -> Maybe (TypeBase dim o))
-> TypeBase dim o -> Maybe (TypeBase dim o)
forall a b. (a -> b) -> a -> b
$ (NoMode -> o) -> TypeBase dim NoMode -> TypeBase dim o
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (o -> NoMode -> o
forall a b. a -> b -> a
const o
o) (ScalarTypeBase dim NoMode -> TypeBase dim NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase dim NoMode
t)
  | Bool
otherwise =
      o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o (Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o)
-> Maybe (Shape dim)
-> Maybe (ScalarTypeBase dim NoMode -> TypeBase dim o)
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 NoMode -> TypeBase dim o)
-> Maybe (ScalarTypeBase dim NoMode) -> Maybe (TypeBase dim o)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ScalarTypeBase dim NoMode -> Maybe (ScalarTypeBase dim NoMode)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase dim NoMode
t
peelArray Int
_ TypeBase dim o
_ = Maybe (TypeBase dim o)
forall a. Maybe a
Nothing

-- | @arrayOf o 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.
arrayOf ::
  Shape dim ->
  TypeBase dim NoMode ->
  TypeBase dim NoMode
arrayOf :: forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf = NoMode -> Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
forall o dim o'.
o -> Shape dim -> TypeBase dim o' -> TypeBase dim o
arrayOfWithAliases NoMode
forall a. Monoid a => a
mempty

-- | Like 'arrayOf', but you can pass in the annotation of the
-- resulting array.
arrayOfWithAliases ::
  o ->
  Shape dim ->
  TypeBase dim o' ->
  TypeBase dim o
arrayOfWithAliases :: forall o dim o'.
o -> Shape dim -> TypeBase dim o' -> TypeBase dim o
arrayOfWithAliases o
o Shape dim
shape2 (Array o'
_ Shape dim
shape1 ScalarTypeBase dim NoMode
et) =
  o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o (Shape dim
shape2 Shape dim -> Shape dim -> Shape dim
forall a. Semigroup a => a -> a -> a
<> Shape dim
shape1) ScalarTypeBase dim NoMode
et
arrayOfWithAliases o
o Shape dim
shape (Scalar ScalarTypeBase dim o'
t) =
  if Shape dim -> Int
forall dim. Shape dim -> Int
shapeRank Shape dim
shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then ScalarTypeBase dim o' -> TypeBase dim o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase dim o'
t TypeBase dim o' -> o -> TypeBase dim o
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`setMode` o
o
    else o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape dim
shape ((o' -> NoMode)
-> ScalarTypeBase dim o' -> ScalarTypeBase dim NoMode
forall b c a. (b -> c) -> ScalarTypeBase a b -> ScalarTypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (NoMode -> o' -> NoMode
forall a b. a -> b -> a
const NoMode
forall a. Monoid a => a
mempty) ScalarTypeBase dim o'
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 o -> TypeBase dim o
stripArray :: forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
n (Array o
o Shape dim
shape ScalarTypeBase dim NoMode
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 =
      o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape dim
shape' ScalarTypeBase dim NoMode
et
  | Bool
otherwise =
      (NoMode -> o) -> TypeBase dim NoMode -> TypeBase dim o
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (o -> NoMode -> o
forall a b. a -> b -> a
const o
o) (ScalarTypeBase dim NoMode -> TypeBase dim NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase dim NoMode
et)
stripArray Int
_ TypeBase dim o
t = TypeBase dim o
t

-- | Create a record type corresponding to a tuple with the given
-- element types.
tupleRecord :: [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord :: forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord = Map Name (TypeBase dim o) -> ScalarTypeBase dim o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase dim o) -> ScalarTypeBase dim o)
-> ([TypeBase dim o] -> Map Name (TypeBase dim o))
-> [TypeBase dim o]
-> ScalarTypeBase dim o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, TypeBase dim o)] -> Map Name (TypeBase dim o)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase dim o)] -> Map Name (TypeBase dim o))
-> ([TypeBase dim o] -> [(Name, TypeBase dim o)])
-> [TypeBase dim o]
-> Map Name (TypeBase dim o)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Name] -> [TypeBase dim o] -> [(Name, TypeBase dim o)]
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 o -> Maybe [TypeBase dim o]
isTupleRecord :: forall dim o. TypeBase dim o -> Maybe [TypeBase dim o]
isTupleRecord (Scalar (Record Map Name (TypeBase dim o)
fs)) = Map Name (TypeBase dim o) -> Maybe [TypeBase dim o]
forall a. Map Name a -> Maybe [a]
areTupleFields Map Name (TypeBase dim o)
fs
isTupleRecord TypeBase dim o
_ = Maybe [TypeBase dim o]
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

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

-- | A special expression representing no known size, but encoding an
-- equivalence class (represented by the integer) that can be used to detect
-- unknown-but-equal sizes. This is important so that we do not throw away size
-- equalities just because the names become unknown to us (e.g. see #2326). The
-- type checker should _never_ produce anySizes - they are a (hopefully
-- temporary) thing introduced by defunctorisation and monomorphisation. They
-- represent a flaw in our implementation. When they occur in a return type,
-- they can be replaced with freshly created existential sizes. When they occur
-- in parameter types, they can be replaced with size parameters.
anySize :: Int -> Size
anySize :: Int -> Exp
anySize Int
x =
  -- The definition here is weird to avoid seeing this as a free variable.
  [Word8] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [Word8] -> SrcLoc -> ExpBase f vn
StringLit ((Char -> Word8) -> String -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> (Char -> Int) -> Char -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
ord) (Int -> String
forall a. Show a => a -> String
show Int
x)) SrcLoc
forall a. Monoid a => a
mempty

-- | If this is any size, retrieve the key for the equivalence class.
isAnySize :: Size -> Maybe Int
isAnySize :: Exp -> Maybe Int
isAnySize (StringLit [Word8]
xs SrcLoc
_) = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ String -> Int
forall a. Read a => String -> a
read (String -> Int) -> String -> Int
forall a b. (a -> b) -> a -> b
$ (Word8 -> Char) -> [Word8] -> String
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Char
chr (Int -> Char) -> (Word8 -> Int) -> Word8 -> Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral) [Word8]
xs
isAnySize Exp
_ = Maybe Int
forall a. Maybe a
Nothing

-- | Override the location of an 'Apply' expression.
--
-- This is useful because 'mkApply' sets the location to be the span of the
-- function and all arguments, but during many frontend transformations we
-- insert additional arguments for which we do not care about the original
-- source location. Our goal is to maintain a correspondence between the
-- original source code and the generated code.
setApplyLoc :: SrcLoc -> Exp -> Exp
setApplyLoc :: SrcLoc -> Exp -> Exp
setApplyLoc SrcLoc
loc (AppExp (Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
_) Info AppRes
info) = AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> NonEmpty (Info (Maybe VName), Exp)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
loc) Info AppRes
info
setApplyLoc SrcLoc
_ Exp
e = Exp
e

-- | 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 o m d1 d2.
  (Monoid o, Monad m) =>
  ([VName] -> d1 -> d2 -> m d1) ->
  TypeBase d1 o ->
  TypeBase d2 o ->
  m (TypeBase d1 o)
matchDims :: forall o (m :: * -> *) d1 d2.
(Monoid o, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 o -> TypeBase d2 o -> m (TypeBase d1 o)
matchDims [VName] -> d1 -> d2 -> m d1
onDims = [VName] -> TypeBase d1 o -> TypeBase d2 o -> m (TypeBase d1 o)
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
forall a. Monoid a => a
mempty
  where
    matchDims' ::
      forall o'. (Monoid o') => [VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
    matchDims' :: forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound TypeBase d1 o'
t1 TypeBase d2 o'
t2 =
      case (TypeBase d1 o'
t1, TypeBase d2 o'
t2) of
        (Array o'
u1 Shape d1
shape1 ScalarTypeBase d1 NoMode
et1, Array o'
u2 Shape d2
shape2 ScalarTypeBase d2 NoMode
et2) ->
          o' -> Shape d1 -> TypeBase d1 o' -> TypeBase d1 o'
forall o dim o'.
o -> Shape dim -> TypeBase dim o' -> TypeBase dim o
arrayOfWithAliases o'
u1
            (Shape d1 -> TypeBase d1 o' -> TypeBase d1 o')
-> m (Shape d1) -> m (TypeBase d1 o' -> TypeBase d1 o')
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 o' -> TypeBase d1 o')
-> m (TypeBase d1 o') -> m (TypeBase d1 o')
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound ((NoMode -> o') -> TypeBase d1 NoMode -> TypeBase d1 o'
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (o' -> NoMode -> o'
forall a b. a -> b -> a
const o'
u2) (ScalarTypeBase d1 NoMode -> TypeBase d1 NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase d1 NoMode
et1)) ((NoMode -> o') -> TypeBase d2 NoMode -> TypeBase d2 o'
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (o' -> NoMode -> o'
forall a b. a -> b -> a
const o'
u2) (ScalarTypeBase d2 NoMode -> TypeBase d2 NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase d2 NoMode
et2))
        (Scalar (Record Map Name (TypeBase d1 o')
f1), Scalar (Record Map Name (TypeBase d2 o')
f2)) ->
          ScalarTypeBase d1 o' -> TypeBase d1 o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase d1 o' -> TypeBase d1 o')
-> (Map Name (TypeBase d1 o') -> ScalarTypeBase d1 o')
-> Map Name (TypeBase d1 o')
-> TypeBase d1 o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase d1 o') -> ScalarTypeBase d1 o'
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record
            (Map Name (TypeBase d1 o') -> TypeBase d1 o')
-> m (Map Name (TypeBase d1 o')) -> m (TypeBase d1 o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((TypeBase d1 o', TypeBase d2 o') -> m (TypeBase d1 o'))
-> Map Name (TypeBase d1 o', TypeBase d2 o')
-> m (Map Name (TypeBase d1 o'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse ((TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o'))
-> (TypeBase d1 o', TypeBase d2 o') -> m (TypeBase d1 o')
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ([VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound)) ((TypeBase d1 o'
 -> TypeBase d2 o' -> (TypeBase d1 o', TypeBase d2 o'))
-> Map Name (TypeBase d1 o')
-> Map Name (TypeBase d2 o')
-> Map Name (TypeBase d1 o', TypeBase d2 o')
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 o')
f1 Map Name (TypeBase d2 o')
f2)
        (Scalar (Sum Map Name [TypeBase d1 o']
cs1), Scalar (Sum Map Name [TypeBase d2 o']
cs2)) ->
          ScalarTypeBase d1 o' -> TypeBase d1 o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase d1 o' -> TypeBase d1 o')
-> (Map Name [TypeBase d1 o'] -> ScalarTypeBase d1 o')
-> Map Name [TypeBase d1 o']
-> TypeBase d1 o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase d1 o'] -> ScalarTypeBase d1 o'
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum
            (Map Name [TypeBase d1 o'] -> TypeBase d1 o')
-> m (Map Name [TypeBase d1 o']) -> m (TypeBase d1 o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(TypeBase d1 o', TypeBase d2 o')] -> m [TypeBase d1 o'])
-> Map Name [(TypeBase d1 o', TypeBase d2 o')]
-> m (Map Name [TypeBase d1 o'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse
              (((TypeBase d1 o', TypeBase d2 o') -> m (TypeBase d1 o'))
-> [(TypeBase d1 o', TypeBase d2 o')] -> m [TypeBase d1 o']
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o'))
-> (TypeBase d1 o', TypeBase d2 o') -> m (TypeBase d1 o')
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ([VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound)))
              (([TypeBase d1 o']
 -> [TypeBase d2 o'] -> [(TypeBase d1 o', TypeBase d2 o')])
-> Map Name [TypeBase d1 o']
-> Map Name [TypeBase d2 o']
-> Map Name [(TypeBase d1 o', TypeBase d2 o')]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith [TypeBase d1 o']
-> [TypeBase d2 o'] -> [(TypeBase d1 o', TypeBase d2 o')]
forall a b. [a] -> [b] -> [(a, b)]
zip Map Name [TypeBase d1 o']
cs1 Map Name [TypeBase d2 o']
cs2)
        ( Scalar (Arrow o'
als1 PName
p1 Diet
d1 TypeBase d1 NoMode
a1 (RetType [VName]
dims1 TypeBase d1 Freshness
b1)),
          Scalar (Arrow o'
als2 PName
p2 Diet
_d2 TypeBase d2 NoMode
a2 (RetType [VName]
dims2 TypeBase d2 Freshness
b2))
          ) ->
            let bound' :: [VName]
bound' = (PName -> Maybe VName) -> [PName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe PName -> Maybe VName
paramName [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 o' -> TypeBase d1 o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase d1 o' -> TypeBase d1 o')
-> m (ScalarTypeBase d1 o') -> m (TypeBase d1 o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( o'
-> PName
-> Diet
-> TypeBase d1 NoMode
-> RetTypeBase d1 Freshness
-> ScalarTypeBase d1 o'
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow (o'
als1 o' -> o' -> o'
forall a. Semigroup a => a -> a -> a
<> o'
als2) PName
p1 Diet
d1
                          (TypeBase d1 NoMode
 -> RetTypeBase d1 Freshness -> ScalarTypeBase d1 o')
-> m (TypeBase d1 NoMode)
-> m (RetTypeBase d1 Freshness -> ScalarTypeBase d1 o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName]
-> TypeBase d1 NoMode
-> TypeBase d2 NoMode
-> m (TypeBase d1 NoMode)
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound' TypeBase d1 NoMode
a1 TypeBase d2 NoMode
a2
                          m (RetTypeBase d1 Freshness -> ScalarTypeBase d1 o')
-> m (RetTypeBase d1 Freshness) -> m (ScalarTypeBase d1 o')
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([VName] -> TypeBase d1 Freshness -> RetTypeBase d1 Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims1 (TypeBase d1 Freshness -> RetTypeBase d1 Freshness)
-> m (TypeBase d1 Freshness) -> m (RetTypeBase d1 Freshness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName]
-> TypeBase d1 Freshness
-> TypeBase d2 Freshness
-> m (TypeBase d1 Freshness)
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound' TypeBase d1 Freshness
b1 TypeBase d2 Freshness
b2)
                      )
        ( Scalar (TypeVar o'
als1 QualName VName
v [TypeArg d1]
targs1),
          Scalar (TypeVar o'
als2 QualName VName
_ [TypeArg d2]
targs2)
          ) ->
            ScalarTypeBase d1 o' -> TypeBase d1 o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase d1 o' -> TypeBase d1 o')
-> ([TypeArg d1] -> ScalarTypeBase d1 o')
-> [TypeArg d1]
-> TypeBase d1 o'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. o' -> QualName VName -> [TypeArg d1] -> ScalarTypeBase d1 o'
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar (o'
als1 o' -> o' -> o'
forall a. Semigroup a => a -> a -> a
<> o'
als2) QualName VName
v
              ([TypeArg d1] -> TypeBase d1 o')
-> m [TypeArg d1] -> m (TypeBase d1 o')
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 o', TypeBase d2 o')
_ -> TypeBase d1 o' -> m (TypeBase d1 o')
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase d1 o'
t1

    matchTypeArg :: [VName] -> TypeArg d1 -> TypeArg d2 -> m (TypeArg d1)
matchTypeArg [VName]
bound (TypeArgType TypeBase d1 NoMode
t1) (TypeArgType TypeBase d2 NoMode
t2) =
      TypeBase d1 NoMode -> TypeArg d1
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (TypeBase d1 NoMode -> TypeArg d1)
-> m (TypeBase d1 NoMode) -> m (TypeArg d1)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName]
-> TypeBase d1 NoMode
-> TypeBase d2 NoMode
-> m (TypeBase d1 NoMode)
forall o'.
Monoid o' =>
[VName] -> TypeBase d1 o' -> TypeBase d2 o' -> m (TypeBase d1 o')
matchDims' [VName]
bound TypeBase d1 NoMode
t1 TypeBase d2 NoMode
t2
    matchTypeArg [VName]
bound (TypeArgDim d1
x) (TypeArgDim d2
y) =
      d1 -> TypeArg d1
forall dim. dim -> TypeArg dim
TypeArgDim (d1 -> TypeArg d1) -> m d1 -> m (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
    matchTypeArg [VName]
_ TypeArg d1
a TypeArg d2
_ = TypeArg d1 -> m (TypeArg d1)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeArg d1
a

    onShapes :: [VName] -> Shape d1 -> Shape d2 -> m (Shape d1)
onShapes [VName]
bound Shape d1
shape1 Shape d2
shape2 =
      [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 annotation of a type - that is, its 'Diet' or 'Freshness',
-- or nothing at all.  If the type is a record or sum type, the
-- annotation of its components will be modified.
setMode :: TypeBase dim o1 -> o2 -> TypeBase dim o2
setMode :: forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
setMode TypeBase dim o1
t o2
o = (o1 -> o2) -> TypeBase dim o1 -> TypeBase dim o2
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (o2 -> o1 -> o2
forall a b. a -> b -> a
const o2
o) TypeBase dim o1
t

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 an Futhark term.  The aliasing will refer to itself, if
-- the term is a non-tuple-typed variable.
typeOf :: ExpBase Info VName -> StructType
typeOf :: Exp -> StructType
typeOf (Literal PrimValue
val SrcLoc
_) = ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
typeOf (IntLit Integer
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (FloatLit Double
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Parens Exp
e SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (TupLit [Exp]
es SrcLoc
_) = ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ [StructType] -> ScalarTypeBase Exp NoMode
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord ([StructType] -> ScalarTypeBase Exp NoMode)
-> [StructType] -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ (Exp -> StructType) -> [Exp] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> StructType
typeOf [Exp]
es
typeOf (RecordLit [FieldBase Info VName]
fs SrcLoc
_) =
  ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Exp NoMode
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name StructType -> ScalarTypeBase Exp NoMode)
-> Map Name StructType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ [(Name, StructType)] -> Map Name StructType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, StructType)] -> Map Name StructType)
-> [(Name, StructType)] -> Map Name StructType
forall a b. (a -> b) -> a -> b
$ (FieldBase Info VName -> (Name, StructType))
-> [FieldBase Info VName] -> [(Name, StructType)]
forall a b. (a -> b) -> [a] -> [b]
map FieldBase Info VName -> (Name, StructType)
record [FieldBase Info VName]
fs
  where
    record :: FieldBase Info VName -> (Name, StructType)
record (RecordFieldExplicit (L Loc
_ Name
name) Exp
e SrcLoc
_) = (Name
name, Exp -> StructType
typeOf Exp
e)
    record (RecordFieldImplicit (L Loc
_ VName
name) (Info StructType
t) SrcLoc
_) = (VName -> Name
baseName VName
name, StructType
t)
typeOf (ArrayLit [Exp]
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  NoMode -> Shape Exp -> ScalarTypeBase Exp NoMode -> StructType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array NoMode
forall a. Monoid a => a
mempty ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Exp
sizeFromInteger ([PrimValue] -> Integer
forall i a. Num i => [a] -> i
genericLength [PrimValue]
vs) SrcLoc
loc]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
t)
typeOf (StringLit [Word8]
vs SrcLoc
loc) =
  NoMode -> Shape Exp -> ScalarTypeBase Exp NoMode -> StructType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array
    NoMode
forall a. Monoid a => a
mempty
    ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [Integer -> SrcLoc -> Exp
sizeFromInteger ([Word8] -> Integer
forall i a. Num i => [a] -> i
genericLength [Word8]
vs) SrcLoc
loc])
    (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (IntType -> PrimType
Unsigned IntType
Int8))
typeOf (Project Name
_ Exp
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Var QualName VName
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Hole (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Ascript Exp
e TypeExp Exp VName
_ SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (Coerce Exp
_ TypeExp Exp VName
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Negate Exp
e SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (Not Exp
e SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (Update Exp
_ [UpdateStep Info VName]
_ Exp
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Assert Exp
_ Exp
e Info Text
_ SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (Lambda [PatBase Info VName ParamType]
params Exp
_ Maybe (TypeExp Exp VName)
_ (Info ResRetType
t) SrcLoc
_) = [PatBase Info VName ParamType] -> ResRetType -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [PatBase Info VName ParamType]
params ResRetType
t
typeOf (OpSection QualName VName
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (OpSectionLeft QualName VName
_ Info StructType
_ Exp
_ (Info (PName, ParamType, Maybe VName)
_, Info (PName
pn, ParamType
pt2)) (Info ResRetType
ret, Info [VName]
_) SrcLoc
_) =
  ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoMode
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow NoMode
forall a. Monoid a => a
mempty PName
pn (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
pt2) (ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
pt2) ResRetType
ret
typeOf (OpSectionRight QualName VName
_ Info StructType
_ Exp
_ (Info (PName
pn, ParamType
pt1), Info (PName, ParamType, Maybe VName)
_) (Info ResRetType
ret) SrcLoc
_) =
  ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoMode
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow NoMode
forall a. Monoid a => a
mempty PName
pn (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
pt1) (ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
pt1) ResRetType
ret
typeOf (UpdateSection [UpdateStep Info VName]
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Constr Name
_ [Exp]
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (Attr AttrInfo VName
_ Exp
e SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (AppExp AppExpBase Info VName
_ (Info AppRes
res)) = AppRes -> StructType
appResType AppRes
res

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

-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes
-- @ts@ as parameters and returns @ret@.
foldFunType :: [TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType :: forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [TypeBase d Diet]
ps RetTypeBase d Freshness
ret =
  let RetType [VName]
_ TypeBase d Freshness
t = (TypeBase d Diet
 -> RetTypeBase d Freshness -> RetTypeBase d Freshness)
-> RetTypeBase d Freshness
-> [TypeBase d Diet]
-> RetTypeBase d Freshness
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TypeBase d Diet
-> RetTypeBase d Freshness -> RetTypeBase d Freshness
forall {dim}.
TypeBase dim Diet
-> RetTypeBase dim Freshness -> RetTypeBase dim Freshness
arrow RetTypeBase d Freshness
ret [TypeBase d Diet]
ps
   in TypeBase d Freshness -> TypeBase d NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase d Freshness
t
  where
    arrow :: TypeBase dim Diet
-> RetTypeBase dim Freshness -> RetTypeBase dim Freshness
arrow TypeBase dim Diet
t1 RetTypeBase dim Freshness
t2 =
      [VName] -> TypeBase dim Freshness -> RetTypeBase dim Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (TypeBase dim Freshness -> RetTypeBase dim Freshness)
-> TypeBase dim Freshness -> RetTypeBase dim Freshness
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim Freshness -> TypeBase dim Freshness
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim Freshness -> TypeBase dim Freshness)
-> ScalarTypeBase dim Freshness -> TypeBase dim Freshness
forall a b. (a -> b) -> a -> b
$ Freshness
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim Freshness
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow Freshness
Nonfresh PName
Unnamed (TypeBase dim Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase dim Diet
t1) (TypeBase dim Diet -> TypeBase dim NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase dim Diet
t1) RetTypeBase dim Freshness
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 o -> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType :: forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType (Scalar (Arrow o
_ PName
p Diet
d TypeBase dim NoMode
t1 (RetType [VName]
_ TypeBase dim Freshness
t2))) =
  let ([(PName, TypeBase dim Diet)]
ps, TypeBase dim NoMode
r) = TypeBase dim Freshness
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType TypeBase dim Freshness
t2
   in ((PName
p, (NoMode -> Diet) -> TypeBase dim NoMode -> TypeBase dim Diet
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Diet -> NoMode -> Diet
forall a b. a -> b -> a
const Diet
d) TypeBase dim NoMode
t1) (PName, TypeBase dim Diet)
-> [(PName, TypeBase dim Diet)] -> [(PName, TypeBase dim Diet)]
forall a. a -> [a] -> [a]
: [(PName, TypeBase dim Diet)]
ps, TypeBase dim NoMode
r)
unfoldFunType TypeBase dim o
t = ([], TypeBase dim o -> TypeBase dim NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase dim o
t)

-- | The type scheme of a value binding, comprising the type
-- parameters and the actual type.
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBindBase Info VName
vb =
  ( ValBindBase Info VName -> [TypeParamBase VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [TypeParamBase vn]
valBindTypeParams ValBindBase Info VName
vb,
    [PatBase Info VName ParamType] -> ResRetType -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType (ValBindBase Info VName -> [PatBase Info VName ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBindBase Info VName
vb) (Info ResRetType -> ResRetType
forall a. Info a -> a
unInfo (ValBindBase Info VName -> Info ResRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f ResRetType
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 ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBindBase Info VName
vb of
      [] -> ResRetType -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims (Info ResRetType -> ResRetType
forall a. Info a -> a
unInfo (ValBindBase Info VName -> Info ResRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f ResRetType
valBindRetType ValBindBase Info VName
vb))
      [PatBase Info VName ParamType]
_ -> []

-- | The qualified type names mentioned in a type.
typeQualVars :: TypeBase dim o -> [QualName VName]
typeQualVars :: forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars TypeBase dim o
t =
  case TypeBase dim o
t of
    Scalar Prim {} -> [QualName VName]
forall a. Monoid a => a
mempty
    Scalar (TypeVar o
_ QualName VName
tn [TypeArg dim]
targs) ->
      QualName VName
tn QualName VName -> [QualName VName] -> [QualName VName]
forall a. a -> [a] -> [a]
: (TypeArg dim -> [QualName VName])
-> [TypeArg dim] -> [QualName VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap TypeArg dim -> [QualName VName]
forall {dim}. TypeArg dim -> [QualName VName]
typeArgFree [TypeArg dim]
targs
    Scalar (Arrow o
_ PName
_ Diet
_ TypeBase dim NoMode
t1 (RetType [VName]
_ TypeBase dim Freshness
t2)) -> TypeBase dim NoMode -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars TypeBase dim NoMode
t1 [QualName VName] -> [QualName VName] -> [QualName VName]
forall a. Semigroup a => a -> a -> a
<> TypeBase dim Freshness -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars TypeBase dim Freshness
t2
    Scalar (Record Map Name (TypeBase dim o)
fields) -> (TypeBase dim o -> [QualName VName])
-> Map Name (TypeBase dim o) -> [QualName VName]
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase dim o -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars Map Name (TypeBase dim o)
fields
    Scalar (Sum Map Name [TypeBase dim o]
cs) -> [[QualName VName]] -> [QualName VName]
forall a. Monoid a => [a] -> a
mconcat ([[QualName VName]] -> [QualName VName])
-> [[QualName VName]] -> [QualName VName]
forall a b. (a -> b) -> a -> b
$ (([TypeBase dim o] -> [[QualName VName]])
-> Map Name [TypeBase dim o] -> [[QualName VName]]
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (([TypeBase dim o] -> [[QualName VName]])
 -> Map Name [TypeBase dim o] -> [[QualName VName]])
-> ((TypeBase dim o -> [QualName VName])
    -> [TypeBase dim o] -> [[QualName VName]])
-> (TypeBase dim o -> [QualName VName])
-> Map Name [TypeBase dim o]
-> [[QualName VName]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase dim o -> [QualName VName])
-> [TypeBase dim o] -> [[QualName VName]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) TypeBase dim o -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars Map Name [TypeBase dim o]
cs
    Array o
_ Shape dim
_ ScalarTypeBase dim NoMode
rt -> TypeBase dim NoMode -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars (TypeBase dim NoMode -> [QualName VName])
-> TypeBase dim NoMode -> [QualName VName]
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim NoMode -> TypeBase dim NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase dim NoMode
rt
  where
    typeArgFree :: TypeArg dim -> [QualName VName]
typeArgFree (TypeArgType TypeBase dim NoMode
ta) = TypeBase dim NoMode -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars TypeBase dim NoMode
ta
    typeArgFree TypeArgDim {} = [QualName VName]
forall a. Monoid a => a
mempty

-- | The type names mentioned in a type.
typeVars :: TypeBase dim o -> S.Set VName
typeVars :: forall dim o. TypeBase dim o -> Set VName
typeVars = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName)
-> (TypeBase dim o -> [VName]) -> TypeBase dim o -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> VName) -> [QualName VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf ([QualName VName] -> [VName])
-> (TypeBase dim o -> [QualName VName])
-> TypeBase dim o
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase dim o -> [QualName VName]
forall dim o. TypeBase dim o -> [QualName VName]
typeQualVars

-- | @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 o -> Bool
orderZero :: forall dim o. TypeBase dim o -> Bool
orderZero Array {} = Bool
True
orderZero (Scalar (Prim PrimType
_)) = Bool
True
orderZero (Scalar (Record Map Name (TypeBase dim o)
fs)) = (TypeBase dim o -> Bool) -> [TypeBase dim o] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim o -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero ([TypeBase dim o] -> Bool) -> [TypeBase dim o] -> Bool
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim o) -> [TypeBase dim o]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase dim o)
fs
orderZero (Scalar TypeVar {}) = Bool
True
orderZero (Scalar Arrow {}) = Bool
False
orderZero (Scalar (Sum Map Name [TypeBase dim o]
cs)) = ([TypeBase dim o] -> Bool) -> Map Name [TypeBase dim o] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((TypeBase dim o -> Bool) -> [TypeBase dim o] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase dim o -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero) Map Name [TypeBase dim o]
cs

-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern
-- have order 0.
patternOrderZero :: Pat (TypeBase d o) -> Bool
patternOrderZero :: forall d o. Pat (TypeBase d o) -> Bool
patternOrderZero = TypeBase d o -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero (TypeBase d o -> Bool)
-> (Pat (TypeBase d o) -> TypeBase d o)
-> Pat (TypeBase d o)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (TypeBase d o) -> TypeBase d o
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType

-- | The set of identifiers bound in a pattern.
patIdents :: PatBase f vn t -> [IdentBase f vn t]
patIdents :: forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents (Id vn
v f t
t SrcLoc
loc) = [vn -> f t -> SrcLoc -> IdentBase f vn t
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident vn
v f t
t SrcLoc
loc]
patIdents (PatParens PatBase f vn t
p SrcLoc
_) = PatBase f vn t -> [IdentBase f vn t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents PatBase f vn t
p
patIdents (TuplePat [PatBase f vn t]
pats SrcLoc
_) = (PatBase f vn t -> [IdentBase f vn t])
-> [PatBase f vn t] -> [IdentBase f vn t]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase f vn t -> [IdentBase f vn t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents [PatBase f vn t]
pats
patIdents (RecordPat [(L Name, PatBase f vn t)]
fs SrcLoc
_) = ((L Name, PatBase f vn t) -> [IdentBase f vn t])
-> [(L Name, PatBase f vn t)] -> [IdentBase f vn t]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatBase f vn t -> [IdentBase f vn t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents (PatBase f vn t -> [IdentBase f vn t])
-> ((L Name, PatBase f vn t) -> PatBase f vn t)
-> (L Name, PatBase f vn t)
-> [IdentBase f vn t]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (L Name, PatBase f vn t) -> PatBase f vn t
forall a b. (a, b) -> b
snd) [(L Name, PatBase f vn t)]
fs
patIdents Wildcard {} = [IdentBase f vn t]
forall a. Monoid a => a
mempty
patIdents (PatAscription PatBase f vn t
p TypeExp (ExpBase f vn) vn
_ SrcLoc
_) = PatBase f vn t -> [IdentBase f vn t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents PatBase f vn t
p
patIdents PatLit {} = [IdentBase f vn t]
forall a. Monoid a => a
mempty
patIdents (PatConstr Name
_ f t
_ [PatBase f vn t]
ps SrcLoc
_) = (PatBase f vn t -> [IdentBase f vn t])
-> [PatBase f vn t] -> [IdentBase f vn t]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase f vn t -> [IdentBase f vn t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents [PatBase f vn t]
ps
patIdents (PatAttr AttrInfo vn
_ PatBase f vn t
p SrcLoc
_) = PatBase f vn t -> [IdentBase f vn t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents PatBase f vn t
p

-- | The set of names bound in a pattern.
patNames :: Pat t -> [VName]
patNames :: forall t. Pat t -> [VName]
patNames = ((VName, t) -> VName) -> [(VName, t)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, t) -> VName
forall a b. (a, b) -> a
fst ([(VName, t)] -> [VName])
-> (Pat t -> [(VName, t)]) -> Pat t -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat t -> [(VName, t)]
forall t. Pat t -> [(VName, t)]
patternMap

-- | Each name bound in a pattern alongside its type.
patternMap :: Pat t -> [(VName, t)]
patternMap :: forall t. Pat t -> [(VName, t)]
patternMap = (IdentBase Info VName t -> (VName, t))
-> [IdentBase Info VName t] -> [(VName, t)]
forall a b. (a -> b) -> [a] -> [b]
map IdentBase Info VName t -> (VName, t)
forall {a} {b}. IdentBase Info a b -> (a, b)
f ([IdentBase Info VName t] -> [(VName, t)])
-> (Pat t -> [IdentBase Info VName t]) -> Pat t -> [(VName, t)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat t -> [IdentBase Info VName t]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
patIdents
  where
    f :: IdentBase Info a b -> (a, b)
f (Ident a
v (Info b
t) SrcLoc
_) = (a
v, b
t)

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

-- | The type matched by the pattern, including shape declarations if present.
patternStructType :: Pat (TypeBase Size o) -> StructType
patternStructType :: forall o. Pat (TypeBase Exp o) -> StructType
patternStructType = TypeBase Exp o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (TypeBase Exp o -> StructType)
-> (Pat (TypeBase Exp o) -> TypeBase Exp o)
-> Pat (TypeBase Exp o)
-> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (TypeBase Exp o) -> TypeBase Exp o
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType

-- | When viewed as a function parameter, does this pattern correspond
-- to a named parameter of some type?
patternParam :: Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam :: forall d. Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam (PatParens PatBase Info VName (TypeBase d Diet)
p SrcLoc
_) =
  PatBase Info VName (TypeBase d Diet)
-> (PName, Diet, TypeBase d NoMode)
forall d. Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam PatBase Info VName (TypeBase d Diet)
p
patternParam (PatAttr AttrInfo VName
_ PatBase Info VName (TypeBase d Diet)
p SrcLoc
_) =
  PatBase Info VName (TypeBase d Diet)
-> (PName, Diet, TypeBase d NoMode)
forall d. Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoMode)
patternParam PatBase Info VName (TypeBase d Diet)
p
patternParam (PatAscription (Id VName
v (Info TypeBase d Diet
t) SrcLoc
_) TypeExp Exp VName
_ SrcLoc
_) =
  (VName -> PName
Named VName
v, TypeBase d Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase d Diet
t, TypeBase d Diet -> TypeBase d NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase d Diet
t)
patternParam (Id VName
v (Info TypeBase d Diet
t) SrcLoc
_) =
  (VName -> PName
Named VName
v, TypeBase d Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase d Diet
t, TypeBase d Diet -> TypeBase d NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase d Diet
t)
patternParam PatBase Info VName (TypeBase d Diet)
p =
  (PName
Unnamed, TypeBase d Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase d Diet
p_t, TypeBase d Diet -> TypeBase d NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase d Diet
p_t)
  where
    p_t :: TypeBase d Diet
p_t = PatBase Info VName (TypeBase d Diet) -> TypeBase d Diet
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType PatBase Info VName (TypeBase d Diet)
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
    [ (Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ PrimType -> Text
forall a. Pretty a => a -> Text
prettyText 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] [ParamType] (RetTypeBase Size Freshness)
  | IntrinsicType Liftedness [TypeParamBase VName] StructType
  | IntrinsicEquality -- Special cased.

intrinsicAcc :: (VName, Intrinsic)
intrinsicAcc :: (VName, Intrinsic)
intrinsicAcc =
  ( VName
acc_v,
    Liftedness -> [TypeParamBase VName] -> StructType -> Intrinsic
IntrinsicType Liftedness
SizeLifted [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] (StructType -> Intrinsic) -> StructType -> Intrinsic
forall a b. (a -> b) -> a -> b
$
      ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$
        NoMode
-> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp NoMode
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar NoMode
forall a. Monoid a => a
mempty (VName -> QualName VName
forall v. v -> QualName v
qualName VName
acc_v) [TypeArg Exp
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 NoMode -> TypeArg dim
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (TypeBase dim NoMode -> TypeArg dim)
-> TypeBase dim NoMode -> TypeArg dim
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim NoMode -> TypeBase dim NoMode
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode
-> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim NoMode
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar NoMode
forall a. Monoid a => a
mempty (VName -> QualName VName
forall v. v -> QualName v
qualName VName
t_v) [])

-- | If this type corresponds to the builtin "acc" type, return the
-- type of the underlying array.
isAccType :: TypeBase d o -> Maybe (TypeBase d NoMode)
isAccType :: forall d o. TypeBase d o -> Maybe (TypeBase d NoMode)
isAccType (Scalar (TypeVar o
_ (QualName [] VName
v) [TypeArgType TypeBase d NoMode
t]))
  | VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== (VName, Intrinsic) -> VName
forall a b. (a, b) -> a
fst (VName, Intrinsic)
intrinsicAcc =
      TypeBase d NoMode -> Maybe (TypeBase d NoMode)
forall a. a -> Maybe a
Just TypeBase d NoMode
t
isAccType TypeBase d o
_ = Maybe (TypeBase d NoMode)
forall a. Maybe a
Nothing

-- | Find the 'VName' corresponding to a builtin.  Crashes if that
-- name cannot be found.
intrinsicVar :: Name -> VName
intrinsicVar :: Name -> VName
intrinsicVar Name
v =
  VName -> Maybe VName -> VName
forall a. a -> Maybe a -> a
fromMaybe VName
forall {b}. b
bad (Maybe VName -> VName) -> Maybe VName -> VName
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> Maybe VName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Name
v Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
==) (Name -> Bool) -> (VName -> Name) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Name
baseName) ([VName] -> Maybe VName) -> [VName] -> Maybe VName
forall a b. (a -> b) -> a -> b
$ Map VName Intrinsic -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName Intrinsic
intrinsics
  where
    bad :: b
bad = String -> b
forall a. HasCallStack => String -> a
error (String -> b) -> String -> b
forall a b. (a -> b) -> a -> b
$ String
"findBuiltin: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Name -> String
nameToString Name
v

mkBinOp :: Name -> StructType -> Exp -> Exp -> Exp
mkBinOp :: Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
op StructType
t Exp
x Exp
y =
  AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
    ( (QualName VName, SrcLoc)
-> Info StructType
-> (Exp, Info (Maybe VName))
-> (Exp, Info (Maybe VName))
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
(QualName vn, SrcLoc)
-> f StructType
-> (ExpBase f vn, f (Maybe VName))
-> (ExpBase f vn, f (Maybe VName))
-> SrcLoc
-> AppExpBase f vn
BinOp
        (VName -> QualName VName
forall v. v -> QualName v
qualName (Name -> VName
intrinsicVar Name
op), SrcLoc
forall a. Monoid a => a
mempty)
        (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t)
        (Exp
x, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
        (Exp
y, Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info Maybe VName
forall a. Maybe a
Nothing)
        SrcLoc
forall a. Monoid a => a
mempty
    )
    (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes StructType
t [])

mkAdd, mkMul :: Exp -> Exp -> Exp
mkAdd :: Exp -> Exp -> Exp
mkAdd = Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
"+" (StructType -> Exp -> Exp -> Exp)
-> StructType -> Exp -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
mkMul :: Exp -> Exp -> Exp
mkMul = Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
"*" (StructType -> Exp -> Exp -> Exp)
-> StructType -> Exp -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- | 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
$
      [(VName, Intrinsic)]
primOp
        [(VName, Intrinsic)]
-> [(VName, Intrinsic)] -> [(VName, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (Int -> (Text, Intrinsic) -> (VName, Intrinsic))
-> [Int] -> [(Text, Intrinsic)] -> [(VName, Intrinsic)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          Int -> (Text, Intrinsic) -> (VName, Intrinsic)
forall {b}. Int -> (Text, b) -> (VName, b)
namify
          [Int
intrinsicStart ..]
          ( [ ( Text
"manifest",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a]
                  [ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
forall a. Monoid a => a
mempty]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Fresh
              ),
              ( Text
"flatten",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                  [Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n, VName
m]) (ScalarTypeBase Exp NoMode -> ParamType)
-> ScalarTypeBase Exp NoMode -> ParamType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array
                    Freshness
Nonfresh
                    ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [VName -> Exp
size VName
n Exp -> Exp -> Exp
`mkMul` VName -> Exp
size VName
m])
                    (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty)
              ),
              ( Text
"unflatten",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [VName -> Exp
size VName
n Exp -> Exp -> Exp
`mkMul` VName -> Exp
size VName
m]) (ScalarTypeBase Exp NoMode -> ParamType)
-> ScalarTypeBase Exp NoMode -> ParamType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Nonfresh ([VName] -> Shape Exp
shape [VName
n, VName
m]) (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty)
              ),
              ( Text
"concat",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                  [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [VName -> Exp
size VName
n Exp -> Exp -> Exp
`mkAdd` VName -> Exp
size VName
m]
              ),
              ( Text
"transpose",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                  [Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n, VName
m]]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Nonfresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
n]
              ),
              ( Text
"scatter",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_l]
                  [ Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Consume ([VName] -> Shape Exp
shape [VName
n]) (ScalarTypeBase Exp NoMode -> ParamType)
-> ScalarTypeBase Exp NoMode -> ParamType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (ScalarTypeBase Exp NoMode -> ParamType)
-> ScalarTypeBase Exp NoMode -> ParamType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
n]) (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty)
              ),
              ( Text
"scatter_2d",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_l]
                  [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n, VName
m],
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (Int -> ScalarTypeBase Exp NoMode
forall {dim} {o}. Int -> ScalarTypeBase dim o
tupInt64 Int
2),
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (ScalarTypeBase Exp NoMode -> ParamType)
-> ScalarTypeBase Exp NoMode -> ParamType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n, VName
m]
              ),
              ( Text
"scatter_3d",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l]
                  [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n, VName
m, VName
k],
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (Int -> ScalarTypeBase Exp NoMode
forall {dim} {o}. Int -> ScalarTypeBase dim o
tupInt64 Int
3),
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (ScalarTypeBase Exp NoMode -> ParamType)
-> ScalarTypeBase Exp NoMode -> ParamType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n, VName
m, VName
k]
              ),
              ( Text
"zip",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                  [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe ([VName] -> Shape Exp
shape [VName
n]),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Diet
Observe ([VName] -> Shape Exp
shape [VName
n])
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> StructType -> Shape Exp -> ResType
forall {o} {dim}.
o
-> TypeBase dim NoMode
-> TypeBase dim NoMode
-> Shape dim
-> TypeBase dim o
tuple_array Freshness
Fresh (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
forall a. Monoid a => a
mempty)
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
              ),
              ( Text
"unzip",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                  [Diet -> StructType -> StructType -> Shape Exp -> ParamType
forall {o} {dim}.
o
-> TypeBase dim NoMode
-> TypeBase dim NoMode
-> Shape dim
-> TypeBase dim o
tuple_array Diet
Observe (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
forall a. Monoid a => a
mempty) (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType)
-> ([(Name, ResType)] -> ResType)
-> [(Name, ResType)]
-> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ([(Name, ResType)] -> ScalarTypeBase Exp Freshness)
-> [(Name, ResType)]
-> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name ResType -> ScalarTypeBase Exp Freshness
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name ResType -> ScalarTypeBase Exp Freshness)
-> ([(Name, ResType)] -> Map Name ResType)
-> [(Name, ResType)]
-> ScalarTypeBase Exp Freshness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, ResType)] -> Map Name ResType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                  ([(Name, ResType)] -> ResRetType)
-> [(Name, ResType)] -> ResRetType
forall a b. (a -> b) -> a -> b
$ [Name] -> [ResType] -> [(Name, ResType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n], Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]]
              ),
              ( Text
"hist_1d",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m],
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n]) (Int -> ScalarTypeBase Exp NoMode
forall {dim} {o}. Int -> ScalarTypeBase dim o
tupInt64 Int
1),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe ([VName] -> Shape Exp
shape [VName
n])
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m]
              ),
              ( Text
"hist_2d",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k],
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n]) (Int -> ScalarTypeBase Exp NoMode
forall {dim} {o}. Int -> ScalarTypeBase dim o
tupInt64 Int
2),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe ([VName] -> Shape Exp
shape [VName
n])
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k]
              ),
              ( Text
"hist_3d",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_m, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k, VName
l],
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n]) (Int -> ScalarTypeBase Exp NoMode
forall {dim} {o}. Int -> ScalarTypeBase dim o
tupInt64 Int
3),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe ([VName] -> Shape Exp
shape [VName
n])
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k, VName
l]
              ),
              ( Text
"map",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
              ),
              ( Text
"reduce",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Fresh)
              ),
              ( Text
"reduce_comm",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Fresh))
              ),
              ( Text
"scan",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n])
              ),
              ( Text
"partition",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int32),
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Freshness
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Freshness)
-> PrimType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  ( [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
k] (ResType -> ResRetType)
-> (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness
-> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResRetType)
-> ScalarTypeBase Exp Freshness -> ResRetType
forall a b. (a -> b) -> a -> b
$
                      [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord
                        [ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                          Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
k]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                        ]
                  )
              ),
              ( Text
"acc_write",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
sp_k, TypeParamBase VName
tp_a]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> StructType -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType Diet
Consume (StructType -> ScalarTypeBase Exp Diet)
-> StructType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ NoMode -> StructType
forall {o}. o -> TypeBase Exp o
array_ka NoMode
forall a. Monoid a => a
mempty,
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType Freshness
Fresh (NoMode -> StructType
forall {o}. o -> TypeBase Exp o
array_ka NoMode
forall a. Monoid a => a
mempty)
              ),
              ( Text
"scatter_stream",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_k, TypeParamBase VName
sp_n]
                  [ Diet -> ParamType
forall {o}. o -> TypeBase Exp o
array_ka Diet
Consume,
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> StructType -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType NoMode
forall a. Monoid a => a
mempty (NoMode -> StructType
forall {o}. o -> TypeBase Exp o
array_ka NoMode
forall a. Monoid a => a
mempty))
                      StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`carr` ( ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
forall a. Monoid a => a
mempty)
                                 StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> StructType -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType Freshness
Nonfresh (StructType -> ScalarTypeBase Exp Freshness)
-> StructType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ NoMode -> Shape Exp -> StructType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a NoMode
forall a. Monoid a => a
mempty (Shape Exp -> StructType) -> Shape Exp -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k])
                             ),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> ResType
forall {o}. o -> TypeBase Exp o
array_ka Freshness
Fresh
              ),
              ( Text
"hist_stream",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_k, TypeParamBase VName
sp_n]
                  [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k],
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe,
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> StructType -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType NoMode
forall a. Monoid a => a
mempty (StructType -> ScalarTypeBase Exp NoMode)
-> StructType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ NoMode -> StructType
forall {o}. o -> TypeBase Exp o
array_ka NoMode
forall a. Monoid a => a
mempty)
                      StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`carr` ( ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
forall a. Monoid a => a
mempty)
                                 StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> StructType -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType Freshness
Nonfresh (StructType -> ScalarTypeBase Exp Freshness)
-> StructType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ NoMode -> Shape Exp -> StructType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a NoMode
forall a. Monoid a => a
mempty (Shape Exp -> StructType) -> Shape Exp -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k])
                             ),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                  (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k]
              ),
              ( Text
"jvp2",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe)
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord [ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Fresh, ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Fresh]
              ),
              ( Text
"vjp2",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Diet
Observe)
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord [ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Fresh, ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Fresh]
              ),
              ( Text
"jmp2",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord
                    [ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Fresh,
                      Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                    ]
              ),
              ( Text
"mjp2",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe),
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                  (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord
                    [ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Fresh,
                      Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                    ]
              ),
              ( Text
"with_vjp",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                  [ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
NoMode) StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh),
                    ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
NoMode)
                      StructType -> ResType -> ParamType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ( ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
NoMode)
                                StructType -> ResType -> ResType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
`arr` ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Nonfresh)
                            ),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Diet
Observe)
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b Freshness
Nonfresh)
              ),
              ( Text
"flatmap",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b, TypeParamBase VName
tp_c, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$
                      Diet
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp Diet
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow Diet
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
NoMode)) (ResRetType -> ScalarTypeBase Exp Diet)
-> ResRetType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$
                        [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
k] (ResType -> ResRetType)
-> ([ResType] -> ResType) -> [ResType] -> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ([ResType] -> ScalarTypeBase Exp Freshness)
-> [ResType]
-> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord ([ResType] -> ResRetType) -> [ResType] -> ResRetType
forall a b. (a -> b) -> a -> b
$
                          [ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Freshness
Nonfresh ([VName] -> Shape Exp
shape [VName
k]),
                            ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_c Freshness
Nonfresh
                          ],
                    Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
m] (ResType -> ResRetType)
-> ([ResType] -> ResType) -> [ResType] -> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ([ResType] -> ScalarTypeBase Exp Freshness)
-> [ResType]
-> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord
                  ([ResType] -> ResRetType) -> [ResType] -> ResRetType
forall a b. (a -> b) -> a -> b
$ [ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
n]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
m]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
Bool),
                      Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
n]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_b Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m],
                      Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_c Freshness
Fresh (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                    ]
              )
            ]
              [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++
              -- Interpreter-only IO builtins. It is somewhat ugly (and
              -- non-extensible) that these have to baked-in here. The alternative
              -- is a general "intrinsic" language construct.
              [ ( Text
"io_loadbytes",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
sp_k]
                    [Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
k]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
Int8)]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
n]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
n]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
Int8)
                ),
                ( Text
"io_loadimg",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
sp_k]
                    [Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
k]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
Int8)]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
n, VName
m]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
n, VName
m]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
Int32)
                ),
                ( Text
"io_loadaudio",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
sp_k]
                    [Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
k]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
Int8)]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
n, VName
m]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ScalarTypeBase Exp NoMode -> ResType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Freshness
Fresh ([VName] -> Shape Exp
shape [VName
n, VName
m]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
FloatType FloatType
Float64)
                ),
                ( Text
"io_loadvalue",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
sp_k, TypeParamBase VName
tp_a]
                    [Diet -> Shape Exp -> ScalarTypeBase Exp NoMode -> ParamType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
k]) (PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp NoMode)
-> PrimType -> ScalarTypeBase Exp NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Unsigned IntType
Int8)]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar
                    (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Freshness -> ScalarTypeBase Exp Freshness
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a Freshness
Fresh
                )
              ]
              [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++
              -- Experimental LMAD ones.
              [ ( Text
"flat_index_2d",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                    [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                    ]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
m, VName
k]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Nonfresh
                    (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k]
                ),
                ( Text
"flat_update_2d",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l]
                    [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k, VName
l]
                    ]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                    (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                ),
                ( Text
"flat_index_3d",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                    [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                    ]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
m, VName
k, VName
l]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Nonfresh
                    (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k, VName
l]
                ),
                ( Text
"flat_update_3d",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n, TypeParamBase VName
sp_k, TypeParamBase VName
sp_l, TypeParamBase VName
sp_p]
                    [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k, VName
l, VName
p]
                    ]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                    (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                ),
                ( Text
"flat_index_4d",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                    [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                    [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64)
                    ]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName
m, VName
k, VName
l, VName
p]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Nonfresh
                    (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
m, VName
k, VName
l, VName
p]
                ),
                ( Text
"flat_update_4d",
                  [TypeParamBase VName] -> [ParamType] -> ResRetType -> 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]
                    [ Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Consume (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      ScalarTypeBase Exp Diet -> ParamType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                      Diet -> Shape Exp -> ParamType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Diet
Observe (Shape Exp -> ParamType) -> Shape Exp -> ParamType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
k, VName
l, VName
p, VName
q]
                    ]
                    (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> Shape Exp -> ResType
forall {o} {dim}. o -> Shape dim -> TypeBase dim o
array_a Freshness
Fresh
                    (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n]
                )
              ]
          )
  where
    primOp :: [(VName, Intrinsic)]
primOp =
      (Int -> (Text, Intrinsic) -> (VName, Intrinsic))
-> [Int] -> [(Text, Intrinsic)] -> [(VName, Intrinsic)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> (Text, Intrinsic) -> (VName, Intrinsic)
forall {b}. Int -> (Text, b) -> (VName, b)
namify [Int
20 ..] ([(Text, Intrinsic)] -> [(VName, Intrinsic)])
-> [(Text, Intrinsic)] -> [(VName, Intrinsic)]
forall a b. (a -> b) -> a -> b
$
        ((Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
 -> (Text, Intrinsic))
-> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
-> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map (Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
-> (Text, Intrinsic)
forall {a} {c}. (a, ([PrimType], PrimType, c)) -> (a, Intrinsic)
primFun (Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> [(Text, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
forall k a. Map k a -> [(k, a)]
M.toList Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
Primitive.primFuns)
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (UnOp -> (Text, Intrinsic)) -> [UnOp] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map UnOp -> (Text, Intrinsic)
unOpFun [UnOp]
Primitive.allUnOps
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (BinOp -> (Text, Intrinsic)) -> [BinOp] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map BinOp -> (Text, Intrinsic)
binOpFun [BinOp]
Primitive.allBinOps
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (CmpOp -> (Text, Intrinsic)) -> [CmpOp] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map CmpOp -> (Text, Intrinsic)
cmpOpFun [CmpOp]
Primitive.allCmpOps
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (ConvOp -> (Text, Intrinsic)) -> [ConvOp] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map ConvOp -> (Text, Intrinsic)
convOpFun [ConvOp]
Primitive.allConvOps
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (IntType -> (Text, Intrinsic)) -> [IntType] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> (Text, Intrinsic)
signFun [IntType]
Primitive.allIntTypes
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (IntType -> (Text, Intrinsic)) -> [IntType] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> (Text, Intrinsic)
unsignFun [IntType]
Primitive.allIntTypes
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++ (PrimType -> (Text, Intrinsic))
-> [PrimType] -> [(Text, Intrinsic)]
forall a b. (a -> b) -> [a] -> [b]
map
            PrimType -> (Text, 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]
            )
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, Intrinsic)]
forall a. [a] -> [a] -> [a]
++
          -- This overrides the ! from Primitive.
          [ ( Text
"!",
              [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
            ),
            ( Text
"neg",
              [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]
++ (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]
                )
                [Maybe PrimType
forall a. Maybe a
Nothing]
                Maybe PrimType
forall a. Maybe a
Nothing
            )
          ]
          [(Text, Intrinsic)] -> [(Text, Intrinsic)] -> [(Text, 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 (Text, Intrinsic))
-> [BinOp] -> [(Text, Intrinsic)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe BinOp -> Maybe (Text, Intrinsic)
mkIntrinsicBinOp [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound]

    intrinsicStart :: Int
intrinsicStart = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ VName -> Int
baseTag ((VName, Intrinsic) -> VName
forall a b. (a, b) -> a
fst ((VName, Intrinsic) -> VName) -> (VName, Intrinsic) -> VName
forall a b. (a -> b) -> a -> b
$ [(VName, Intrinsic)] -> (VName, Intrinsic)
forall a. HasCallStack => [a] -> a
last [(VName, Intrinsic)]
primOp)

    [VName
a, VName
b, VName
c, 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 ((Text -> Name) -> [Text] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Name
nameFromText [Text
"a", Text
"b", Text
"c", Text
"n", Text
"m", Text
"k", Text
"l", Text
"p", Text
"q"]) [Int
0 ..]

    t_a :: o -> ScalarTypeBase dim o
t_a o
o = o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o (VName -> QualName VName
forall v. v -> QualName v
qualName VName
a) []
    array_a :: o -> Shape dim -> TypeBase dim o
array_a o
o Shape dim
s = o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape dim
s (ScalarTypeBase dim NoMode -> TypeBase dim o)
-> ScalarTypeBase dim NoMode -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase dim NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty
    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 :: o -> ScalarTypeBase dim o
t_b o
o = o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o (VName -> QualName VName
forall v. v -> QualName v
qualName VName
b) []
    array_b :: o -> Shape dim -> TypeBase dim o
array_b o
o Shape dim
s = o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape dim
s (ScalarTypeBase dim NoMode -> TypeBase dim o)
-> ScalarTypeBase dim NoMode -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase dim NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_b NoMode
forall a. Monoid a => a
mempty
    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

    t_c :: o -> ScalarTypeBase dim o
t_c o
o = o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o (VName -> QualName VName
forall v. v -> QualName v
qualName VName
c) []
    array_c :: o -> Shape dim -> TypeBase dim o
array_c o
o Shape dim
s = o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape dim
s (ScalarTypeBase dim NoMode -> TypeBase dim o)
-> ScalarTypeBase dim NoMode -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase dim NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_c NoMode
forall a. Monoid a => a
mempty
    tp_c :: TypeParamBase VName
tp_c = Liftedness -> VName -> SrcLoc -> TypeParamBase VName
forall vn. Liftedness -> vn -> SrcLoc -> TypeParamBase vn
TypeParamType Liftedness
Unlifted VName
c 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]

    size :: VName -> Exp
size = (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName SrcLoc
forall a. Monoid a => a
mempty (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName
    shape :: [VName] -> Shape Exp
shape = [Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape ([Exp] -> Shape Exp) -> ([VName] -> [Exp]) -> [VName] -> Shape Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Exp) -> [VName] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Exp
size

    tuple_array :: o
-> TypeBase dim NoMode
-> TypeBase dim NoMode
-> Shape dim
-> TypeBase dim o
tuple_array o
o TypeBase dim NoMode
x TypeBase dim NoMode
y Shape dim
s =
      o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape dim
s (Map Name (TypeBase dim NoMode) -> ScalarTypeBase dim NoMode
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record ([(Name, TypeBase dim NoMode)] -> Map Name (TypeBase dim NoMode)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase dim NoMode)] -> Map Name (TypeBase dim NoMode))
-> [(Name, TypeBase dim NoMode)] -> Map Name (TypeBase dim NoMode)
forall a b. (a -> b) -> a -> b
$ [Name] -> [TypeBase dim NoMode] -> [(Name, TypeBase dim NoMode)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [TypeBase dim NoMode
x, TypeBase dim NoMode
y]))

    arr :: TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
arr TypeBase dim NoMode
x TypeBase dim Freshness
y = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe TypeBase dim NoMode
x ([VName] -> TypeBase dim Freshness -> RetTypeBase dim Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase dim Freshness
y)
    carr :: TypeBase dim NoMode -> TypeBase dim Freshness -> TypeBase dim o
carr TypeBase dim NoMode
x TypeBase dim Freshness
y = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
forall a. Monoid a => a
mempty PName
Unnamed Diet
Consume TypeBase dim NoMode
x ([VName] -> TypeBase dim Freshness -> RetTypeBase dim Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase dim Freshness
y)

    array_ka :: o -> TypeBase Exp o
array_ka o
o = o -> Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
k) SrcLoc
forall a. Monoid a => a
mempty]) (ScalarTypeBase Exp NoMode -> TypeBase Exp o)
-> ScalarTypeBase Exp NoMode -> TypeBase Exp o
forall a b. (a -> b) -> a -> b
$ NoMode -> ScalarTypeBase Exp NoMode
forall {o} {dim}. o -> ScalarTypeBase dim o
t_a NoMode
forall a. Monoid a => a
mempty

    accType :: o -> TypeBase dim NoMode -> ScalarTypeBase dim o
accType o
o TypeBase dim NoMode
t =
      o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o (VName -> QualName VName
forall v. v -> QualName v
qualName ((VName, Intrinsic) -> VName
forall a b. (a, b) -> a
fst (VName, Intrinsic)
intrinsicAcc)) [TypeBase dim NoMode -> TypeArg dim
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType TypeBase dim NoMode
t]

    namify :: Int -> (Text, b) -> (VName, b)
namify Int
i (Text
x, b
y) = (Name -> Int -> VName
VName (Text -> Name
nameFromText Text
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 -> (Text, Intrinsic)
unOpFun UnOp
bop = (UnOp -> Text
forall a. Pretty a => a -> Text
prettyText 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 -> (Text, Intrinsic)
binOpFun BinOp
bop = (BinOp -> Text
forall a. Pretty a => a -> Text
prettyText 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 -> (Text, Intrinsic)
cmpOpFun CmpOp
bop = (CmpOp -> Text
forall a. Pretty a => a -> Text
prettyText 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 -> (Text, Intrinsic)
convOpFun ConvOp
cop = (ConvOp -> Text
forall a. Pretty a => a -> Text
prettyText 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 -> (Text, Intrinsic)
signFun IntType
t = (Text
"sign_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> IntType -> Text
forall a. Pretty a => a -> Text
prettyText 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 -> (Text, Intrinsic)
unsignFun IntType
t = (Text
"unsign_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> IntType -> Text
forall a. Pretty a => a -> Text
prettyText 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 -> (Text, Intrinsic)
intrinsicPrim PrimType
t = (PrimType -> Text
forall a. Pretty a => a -> Text
prettyText PrimType
t, Liftedness -> [TypeParamBase VName] -> StructType -> Intrinsic
IntrinsicType Liftedness
Unlifted [] (StructType -> Intrinsic) -> StructType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
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 (T.Text, Intrinsic)
    mkIntrinsicBinOp :: BinOp -> Maybe (Text, Intrinsic)
mkIntrinsicBinOp BinOp
op = do
      op' <- BinOp -> Maybe Intrinsic
intrinsicBinOp BinOp
op
      pure (prettyText op, 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 o
tupInt64 Int
1 =
      PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase dim o)
-> PrimType -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
    tupInt64 Int
x =
      [TypeBase dim o] -> ScalarTypeBase dim o
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord ([TypeBase dim o] -> ScalarTypeBase dim o)
-> [TypeBase dim o] -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ Int -> TypeBase dim o -> [TypeBase dim o]
forall a. Int -> a -> [a]
replicate Int
x (TypeBase dim o -> [TypeBase dim o])
-> TypeBase dim o -> [TypeBase dim o]
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase dim o)
-> PrimType -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- | 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

-- | Is this the name of an intrinsic?
isIntrinsic :: VName -> Bool
isIntrinsic :: VName -> Bool
isIntrinsic = (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag) (Int -> Bool) -> (VName -> Int) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Int
baseTag

-- | 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, Loc)]
progImports :: forall (f :: * -> *) vn. ProgBase f vn -> [(String, Loc)]
progImports = (DecBase f vn -> [(String, Loc)])
-> [DecBase f vn] -> [(String, Loc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, Loc)]
decImports ([DecBase f vn] -> [(String, Loc)])
-> (ProgBase f vn -> [DecBase f vn])
-> ProgBase f vn
-> [(String, Loc)]
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, Loc)]
decImports :: forall (f :: * -> *) vn. DecBase f vn -> [(String, Loc)]
decImports (OpenDec ModExpBase f vn
x SrcLoc
_) = ModExpBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, Loc)]
modExpImports ModExpBase f vn
x
decImports (ModDec ModBindBase f vn
md) = ModExpBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, Loc)]
modExpImports (ModExpBase f vn -> [(String, Loc)])
-> ModExpBase f vn -> [(String, Loc)]
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 ModTypeDec {} = []
decImports TypeDec {} = []
decImports ValDec {} = []
decImports (LocalDec DecBase f vn
d SrcLoc
_) = DecBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, Loc)]
decImports DecBase f vn
d
decImports (ImportDec String
x f ImportName
_ SrcLoc
loc) = [(String
x, SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc)]

modExpImports :: ModExpBase f vn -> [(String, Loc)]
modExpImports :: forall (f :: * -> *) vn. ModExpBase f vn -> [(String, Loc)]
modExpImports ModVar {} = []
modExpImports (ModParens ModExpBase f vn
p SrcLoc
_) = ModExpBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, Loc)]
modExpImports ModExpBase f vn
p
modExpImports (ModImport String
f f ImportName
_ SrcLoc
loc) = [(String
f, SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc)]
modExpImports (ModDecs [DecBase f vn]
ds SrcLoc
_) = (DecBase f vn -> [(String, Loc)])
-> [DecBase f vn] -> [(String, Loc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, Loc)]
decImports [DecBase f vn]
ds
modExpImports (ModApply ModExpBase f vn
_ ModExpBase f vn
me f (Map VName VName)
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, Loc)]
modExpImports ModExpBase f vn
me
modExpImports (ModAscript ModExpBase f vn
me ModTypeExpBase f vn
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(String, Loc)]
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 m a. Monoid m => (a -> m) -> Set a -> m
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 m a. Monoid m => (a -> m) -> Set a -> m
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 m a. Monoid m => (a -> m) -> [a] -> m
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 (ModTypeDec ModTypeBindBase f a
sb) =
          a -> Set a -> Map a (Set a)
forall k a. k -> a -> Map k a
M.singleton (ModTypeBindBase f a -> a
forall (f :: * -> *) vn. ModTypeBindBase f vn -> vn
modTypeName ModTypeBindBase f a
sb) (ModTypeExpBase f a -> Set a
forall {a} {f :: * -> *}. Ord a => ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeBindBase f a -> ModTypeExpBase f a
forall (f :: * -> *) vn.
ModTypeBindBase f vn -> ModTypeExpBase f vn
modTypeExp ModTypeBindBase 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

        onModTypeExp :: ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeVar 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
        onModTypeExp (ModTypeParens ModTypeExpBase f a
e SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp (ModTypeSpecs [SpecBase f a]
ss SrcLoc
_) = (SpecBase f a -> Set a) -> [SpecBase f a] -> Set a
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap SpecBase f a -> Set a
onSpec [SpecBase f a]
ss
        onModTypeExp (ModTypeWith ModTypeExpBase f a
e TypeRefBase f a
_ SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp (ModTypeArrow Maybe a
_ ModTypeExpBase f a
e1 ModTypeExpBase f a
e2 SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase 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 ModTypeExpBase 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
<> ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onSpec (IncludeSpec ModTypeExpBase f a
e SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e

    mtypes_used :: Set VName
mtypes_used = (DecBase Info VName -> Set VName)
-> [DecBase Info VName] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
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
-> ((ModTypeExpBase f VName, f (Map VName VName)) -> Set VName)
-> Maybe (ModTypeExpBase 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 (ModTypeExpBase f VName -> Set VName
forall {a} {f :: * -> *}. Ord a => ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeExpBase f VName -> Set VName)
-> ((ModTypeExpBase f VName, f (Map VName VName))
    -> ModTypeExpBase f VName)
-> (ModTypeExpBase f VName, f (Map VName VName))
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ModTypeExpBase f VName, f (Map VName VName))
-> ModTypeExpBase f VName
forall a b. (a, b) -> a
fst) (ModBindBase f VName
-> Maybe (ModTypeExpBase f VName, f (Map VName VName))
forall (f :: * -> *) vn.
ModBindBase f vn
-> Maybe (ModTypeExpBase f vn, f (Map VName VName))
modType 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 ModTypeDec {} = 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 ModTypeExpBase 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
<> ModTypeExpBase f VName -> Set VName
forall {a} {f :: * -> *}. Ord a => ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f VName
se
        onModExp (ModLambda ModParamBase f VName
p Maybe (ModTypeExpBase 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
-> ((ModTypeExpBase f VName, f (Map VName VName)) -> Set VName)
-> Maybe (ModTypeExpBase 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 (ModTypeExpBase f VName -> Set VName
forall {a} {f :: * -> *}. Ord a => ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeExpBase f VName -> Set VName)
-> ((ModTypeExpBase f VName, f (Map VName VName))
    -> ModTypeExpBase f VName)
-> (ModTypeExpBase f VName, f (Map VName VName))
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ModTypeExpBase f VName, f (Map VName VName))
-> ModTypeExpBase f VName
forall a b. (a, b) -> a
fst) Maybe (ModTypeExpBase 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 = ModTypeExpBase f VName -> Set VName
forall {a} {f :: * -> *}. Ord a => ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeExpBase f VName -> Set VName)
-> (ModParamBase f VName -> ModTypeExpBase f VName)
-> ModParamBase f VName
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModParamBase f VName -> ModTypeExpBase f VName
forall (f :: * -> *) vn. ModParamBase f vn -> ModTypeExpBase f vn
modParamType

        onModTypeExp :: ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeVar 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
        onModTypeExp (ModTypeParens ModTypeExpBase f a
e SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp ModTypeSpecs {} = Set a
forall a. Monoid a => a
mempty
        onModTypeExp (ModTypeWith ModTypeExpBase f a
e TypeRefBase f a
_ SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp (ModTypeArrow Maybe a
_ ModTypeExpBase f a
e1 ModTypeExpBase f a
e2 SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e2

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

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

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

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

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

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

-- | Strip semantically irrelevant stuff from the top level of
-- expression.  This is used to provide a slightly fuzzy notion of
-- expression equality.
--
-- Ideally we'd implement unification on a simpler representation that
-- simply didn't allow us.
stripExp :: Exp -> Maybe Exp
stripExp :: Exp -> Maybe Exp
stripExp (Parens Exp
e SrcLoc
_) = Exp -> Maybe Exp
stripExp Exp
e Maybe Exp -> Maybe Exp -> Maybe Exp
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e
stripExp (Assert Exp
_ Exp
e Info Text
_ SrcLoc
_) = Exp -> Maybe Exp
stripExp Exp
e Maybe Exp -> Maybe Exp -> Maybe Exp
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e
stripExp (Attr AttrInfo VName
_ Exp
e SrcLoc
_) = Exp -> Maybe Exp
stripExp Exp
e Maybe Exp -> Maybe Exp -> Maybe Exp
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e
stripExp (Ascript Exp
e TypeExp Exp VName
_ SrcLoc
_) = Exp -> Maybe Exp
stripExp Exp
e Maybe Exp -> Maybe Exp -> Maybe Exp
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e
stripExp Exp
_ = Maybe Exp
forall a. Maybe a
Nothing

-- | All non-trivial subexpressions (as by stripExp) of some
-- expression, not including the expression itself.
subExps :: Exp -> [Exp]
subExps :: Exp -> [Exp]
subExps Exp
e
  | Just Exp
e' <- Exp -> Maybe Exp
stripExp Exp
e = Exp -> [Exp]
subExps Exp
e'
  | Bool
otherwise = ASTMapper (StateT [Exp] Identity)
-> Exp -> StateT [Exp] Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper (StateT [Exp] Identity)
mapper Exp
e StateT [Exp] Identity Exp -> [Exp] -> [Exp]
forall s a. State s a -> s -> s
`execState` [Exp]
forall a. Monoid a => a
mempty
  where
    mapOnExp :: Exp -> StateT [Exp] Identity Exp
mapOnExp Exp
e'
      | Just Exp
e'' <- Exp -> Maybe Exp
stripExp Exp
e' = Exp -> StateT [Exp] Identity Exp
mapOnExp Exp
e''
      | Bool
otherwise = do
          ([Exp] -> [Exp]) -> StateT [Exp] Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Exp
e' Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
:)
          ASTMapper (StateT [Exp] Identity)
-> Exp -> StateT [Exp] Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper (StateT [Exp] Identity)
mapper Exp
e'
    mapper :: ASTMapper (StateT [Exp] Identity)
mapper = ASTMapper (StateT [Exp] Identity)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp}

similarSlices :: Slice -> Slice -> Maybe [(Exp, Exp)]
similarSlices :: Slice -> Slice -> Maybe [(Exp, Exp)]
similarSlices Slice
slice1 Slice
slice2
  | Slice -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Slice
slice1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Slice -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Slice
slice2 = do
      [[(Exp, Exp)]] -> [(Exp, Exp)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Exp, Exp)]] -> [(Exp, Exp)])
-> Maybe [[(Exp, Exp)]] -> Maybe [(Exp, Exp)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (DimIndexBase Info VName
 -> DimIndexBase Info VName -> Maybe [(Exp, Exp)])
-> Slice -> Slice -> Maybe [[(Exp, Exp)]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM DimIndexBase Info VName
-> DimIndexBase Info VName -> Maybe [(Exp, Exp)]
forall {f :: * -> *} {vn} {f :: * -> *} {vn}.
DimIndexBase f vn
-> DimIndexBase f vn -> Maybe [(ExpBase f vn, ExpBase f vn)]
match Slice
slice1 Slice
slice2
  | Bool
otherwise = Maybe [(Exp, Exp)]
forall a. Maybe a
Nothing
  where
    match :: DimIndexBase f vn
-> DimIndexBase f vn -> Maybe [(ExpBase f vn, ExpBase f vn)]
match (DimFix ExpBase f vn
e1) (DimFix ExpBase f vn
e2) = [(ExpBase f vn, ExpBase f vn)]
-> Maybe [(ExpBase f vn, ExpBase f vn)]
forall a. a -> Maybe a
Just [(ExpBase f vn
e1, ExpBase f vn
e2)]
    match (DimSlice Maybe (ExpBase f vn)
a1 Maybe (ExpBase f vn)
b1 Maybe (ExpBase f vn)
c1) (DimSlice Maybe (ExpBase f vn)
a2 Maybe (ExpBase f vn)
b2 Maybe (ExpBase f vn)
c2) =
      [[(ExpBase f vn, ExpBase f vn)]] -> [(ExpBase f vn, ExpBase f vn)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(ExpBase f vn, ExpBase f vn)]]
 -> [(ExpBase f vn, ExpBase f vn)])
-> Maybe [[(ExpBase f vn, ExpBase f vn)]]
-> Maybe [(ExpBase f vn, ExpBase f vn)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Maybe [(ExpBase f vn, ExpBase f vn)]]
-> Maybe [[(ExpBase f vn, ExpBase f vn)]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [(Maybe (ExpBase f vn), Maybe (ExpBase f vn))
-> Maybe [(ExpBase f vn, ExpBase f vn)]
forall {a} {b}. (Maybe a, Maybe b) -> Maybe [(a, b)]
pair (Maybe (ExpBase f vn)
a1, Maybe (ExpBase f vn)
a2), (Maybe (ExpBase f vn), Maybe (ExpBase f vn))
-> Maybe [(ExpBase f vn, ExpBase f vn)]
forall {a} {b}. (Maybe a, Maybe b) -> Maybe [(a, b)]
pair (Maybe (ExpBase f vn)
b1, Maybe (ExpBase f vn)
b2), (Maybe (ExpBase f vn), Maybe (ExpBase f vn))
-> Maybe [(ExpBase f vn, ExpBase f vn)]
forall {a} {b}. (Maybe a, Maybe b) -> Maybe [(a, b)]
pair (Maybe (ExpBase f vn)
c1, Maybe (ExpBase f vn)
c2)]
    match DimIndexBase f vn
_ DimIndexBase f vn
_ = Maybe [(ExpBase f vn, ExpBase f vn)]
forall a. Maybe a
Nothing
    pair :: (Maybe a, Maybe b) -> Maybe [(a, b)]
pair (Maybe a
Nothing, Maybe b
Nothing) = [(a, b)] -> Maybe [(a, b)]
forall a. a -> Maybe a
Just []
    pair (Just a
x, Just b
y) = [(a, b)] -> Maybe [(a, b)]
forall a. a -> Maybe a
Just [(a
x, b
y)]
    pair (Maybe a, Maybe b)
_ = Maybe [(a, b)]
forall a. Maybe a
Nothing

-- | If these two expressions are structurally similar at top level as
-- sizes, produce their subexpressions (which are not necessarily
-- similar, but you can check for that!). This is the machinery
-- underlying expresssion unification. We assume that the expressions
-- have the same type.
similarExps :: Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps :: Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps Exp
e1 Exp
e2 | Exp -> ExpBase NoInfo VName
bareExp Exp
e1 ExpBase NoInfo VName -> ExpBase NoInfo VName -> Bool
forall a. Eq a => a -> a -> Bool
== Exp -> ExpBase NoInfo VName
bareExp Exp
e2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
similarExps Exp
e1 Exp
e2 | Just Exp
e1' <- Exp -> Maybe Exp
stripExp Exp
e1 = Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps Exp
e1' Exp
e2
similarExps Exp
e1 Exp
e2 | Just Exp
e2' <- Exp -> Maybe Exp
stripExp Exp
e2 = Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps Exp
e1 Exp
e2'
similarExps (IntLit Integer
x Info StructType
_ SrcLoc
_) (Literal PrimValue
v SrcLoc
_) =
  case PrimValue
v of
    SignedValue (Int8Value Int8
y) | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Int8 -> Integer
forall a. Integral a => a -> Integer
toInteger Int8
y -> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
    SignedValue (Int16Value Int16
y) | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Int16 -> Integer
forall a. Integral a => a -> Integer
toInteger Int16
y -> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
    SignedValue (Int32Value Int32
y) | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger Int32
y -> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
    SignedValue (Int64Value Int64
y) | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Integer
forall a. Integral a => a -> Integer
toInteger Int64
y -> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
    PrimValue
_ -> Maybe [(Exp, Exp)]
forall a. Maybe a
Nothing
similarExps
  (AppExp (BinOp (QualName VName
op1, SrcLoc
_) Info StructType
_ (Exp
x1, Info (Maybe VName)
_) (Exp
y1, Info (Maybe VName)
_) SrcLoc
_) Info AppRes
_)
  (AppExp (BinOp (QualName VName
op2, SrcLoc
_) Info StructType
_ (Exp
x2, Info (Maybe VName)
_) (Exp
y2, Info (Maybe VName)
_) SrcLoc
_) Info AppRes
_)
    | QualName VName
op1 QualName VName -> QualName VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName
op2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
x1, Exp
x2), (Exp
y1, Exp
y2)]
similarExps (AppExp (Apply Exp
f1 NonEmpty (Info (Maybe VName), Exp)
args1 SrcLoc
_) Info AppRes
_) (AppExp (Apply Exp
f2 NonEmpty (Info (Maybe VName), Exp)
args2 SrcLoc
_) Info AppRes
_)
  | Exp
f1 Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
f2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just ([(Exp, Exp)] -> Maybe [(Exp, Exp)])
-> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Exp] -> [(Exp, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Info (Maybe VName), Exp) -> Exp)
-> [(Info (Maybe VName), Exp)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> Exp
forall a b. (a, b) -> b
snd ([(Info (Maybe VName), Exp)] -> [Exp])
-> [(Info (Maybe VName), Exp)] -> [Exp]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args1) (((Info (Maybe VName), Exp) -> Exp)
-> [(Info (Maybe VName), Exp)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> Exp
forall a b. (a, b) -> b
snd ([(Info (Maybe VName), Exp)] -> [Exp])
-> [(Info (Maybe VName), Exp)] -> [Exp]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args2)
similarExps (AppExp (Index Exp
arr1 Slice
slice1 SrcLoc
_) Info AppRes
_) (AppExp (Index Exp
arr2 Slice
slice2 SrcLoc
_) Info AppRes
_)
  | Exp
arr1 Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
arr2,
    Slice -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Slice
slice1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Slice -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Slice
slice2 =
      Slice -> Slice -> Maybe [(Exp, Exp)]
similarSlices Slice
slice1 Slice
slice2
similarExps (TupLit [Exp]
es1 SrcLoc
_) (TupLit [Exp]
es2 SrcLoc
_)
  | [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es2 =
      [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just ([(Exp, Exp)] -> Maybe [(Exp, Exp)])
-> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Exp] -> [(Exp, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
es1 [Exp]
es2
similarExps (RecordLit [FieldBase Info VName]
fs1 SrcLoc
_) (RecordLit [FieldBase Info VName]
fs2 SrcLoc
_)
  | [FieldBase Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FieldBase Info VName]
fs1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [FieldBase Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FieldBase Info VName]
fs2 =
      (FieldBase Info VName -> FieldBase Info VName -> Maybe (Exp, Exp))
-> [FieldBase Info VName]
-> [FieldBase Info VName]
-> Maybe [(Exp, Exp)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM FieldBase Info VName -> FieldBase Info VName -> Maybe (Exp, Exp)
forall {f :: * -> *} {vn} {f :: * -> *} {vn}.
FieldBase f vn
-> FieldBase f vn -> Maybe (ExpBase f vn, ExpBase f vn)
onFields [FieldBase Info VName]
fs1 [FieldBase Info VName]
fs2
  where
    onFields :: FieldBase f vn
-> FieldBase f vn -> Maybe (ExpBase f vn, ExpBase f vn)
onFields (RecordFieldExplicit (L Loc
_ Name
n1) ExpBase f vn
fe1 SrcLoc
_) (RecordFieldExplicit (L Loc
_ Name
n2) ExpBase f vn
fe2 SrcLoc
_)
      | Name
n1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n2 = (ExpBase f vn, ExpBase f vn) -> Maybe (ExpBase f vn, ExpBase f vn)
forall a. a -> Maybe a
Just (ExpBase f vn
fe1, ExpBase f vn
fe2)
    onFields (RecordFieldImplicit (L Loc
_ vn
vn1) f StructType
ty1 SrcLoc
_) (RecordFieldImplicit (L Loc
_ vn
vn2) f StructType
ty2 SrcLoc
_) =
      (ExpBase f vn, ExpBase f vn) -> Maybe (ExpBase f vn, ExpBase f vn)
forall a. a -> Maybe a
Just (QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
vn1) f StructType
ty1 SrcLoc
forall a. Monoid a => a
mempty, QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
vn2) f StructType
ty2 SrcLoc
forall a. Monoid a => a
mempty)
    onFields FieldBase f vn
_ FieldBase f vn
_ = Maybe (ExpBase f vn, ExpBase f vn)
forall a. Maybe a
Nothing
similarExps (ArrayLit [Exp]
es1 Info StructType
_ SrcLoc
_) (ArrayLit [Exp]
es2 Info StructType
_ SrcLoc
_)
  | [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es2 =
      [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just ([(Exp, Exp)] -> Maybe [(Exp, Exp)])
-> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Exp] -> [(Exp, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
es1 [Exp]
es2
similarExps (Project Name
field1 Exp
e1 Info StructType
_ SrcLoc
_) (Project Name
field2 Exp
e2 Info StructType
_ SrcLoc
_)
  | Name
field1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
field2 =
      [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
e1, Exp
e2)]
similarExps (Negate Exp
e1 SrcLoc
_) (Negate Exp
e2 SrcLoc
_) =
  [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
e1, Exp
e2)]
similarExps (Not Exp
e1 SrcLoc
_) (Not Exp
e2 SrcLoc
_) =
  [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
e1, Exp
e2)]
similarExps (Constr Name
n1 [Exp]
es1 Info StructType
_ SrcLoc
_) (Constr Name
n2 [Exp]
es2 Info StructType
_ SrcLoc
_)
  | [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es2,
    Name
n1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n2 =
      [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just ([(Exp, Exp)] -> Maybe [(Exp, Exp)])
-> [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Exp] -> [(Exp, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
es1 [Exp]
es2
similarExps (Update Exp
e1 [UpdateStep Info VName]
steps1 Exp
e'1 Info StructType
_ SrcLoc
_) (Update Exp
e2 [UpdateStep Info VName]
steps2 Exp
e'2 Info StructType
_ SrcLoc
_)
  | [UpdateStep Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [UpdateStep Info VName]
steps1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [UpdateStep Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [UpdateStep Info VName]
steps2 = do
      step_pairs <- [[(Exp, Exp)]] -> [(Exp, Exp)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Exp, Exp)]] -> [(Exp, Exp)])
-> Maybe [[(Exp, Exp)]] -> Maybe [(Exp, Exp)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (UpdateStep Info VName
 -> UpdateStep Info VName -> Maybe [(Exp, Exp)])
-> [UpdateStep Info VName]
-> [UpdateStep Info VName]
-> Maybe [[(Exp, Exp)]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM UpdateStep Info VName
-> UpdateStep Info VName -> Maybe [(Exp, Exp)]
similarStep [UpdateStep Info VName]
steps1 [UpdateStep Info VName]
steps2
      pure $ [(e1, e2), (e'1, e'2)] ++ step_pairs
  where
    similarStep :: UpdateStep Info VName
-> UpdateStep Info VName -> Maybe [(Exp, Exp)]
similarStep (UpdateStepField Name
f1) (UpdateStepField Name
f2)
      | Name
f1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
f2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
    similarStep (UpdateStepSlice Slice
s1) (UpdateStepSlice Slice
s2) =
      Slice -> Slice -> Maybe [(Exp, Exp)]
similarSlices Slice
s1 Slice
s2
    similarStep UpdateStep Info VName
_ UpdateStep Info VName
_ = Maybe [(Exp, Exp)]
forall a. Maybe a
Nothing
similarExps (OpSection QualName VName
op1 Info StructType
_ SrcLoc
_) (OpSection QualName VName
op2 Info StructType
_ SrcLoc
_)
  | QualName VName
op1 QualName VName -> QualName VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName
op2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
similarExps (OpSectionLeft QualName VName
op1 Info StructType
_ Exp
x1 (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
_ (Info ResRetType, Info [VName])
_ SrcLoc
_) (OpSectionLeft QualName VName
op2 Info StructType
_ Exp
x2 (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
_ (Info ResRetType, Info [VName])
_ SrcLoc
_)
  | QualName VName
op1 QualName VName -> QualName VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName
op2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
x1, Exp
x2)]
similarExps (OpSectionRight QualName VName
op1 Info StructType
_ Exp
x1 (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
_ Info ResRetType
_ SrcLoc
_) (OpSectionRight QualName VName
op2 Info StructType
_ Exp
x2 (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
_ Info ResRetType
_ SrcLoc
_)
  | QualName VName
op1 QualName VName -> QualName VName -> Bool
forall a. Eq a => a -> a -> Bool
== QualName VName
op2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
x1, Exp
x2)]
similarExps (UpdateSection [UpdateStep Info VName]
steps1 Info StructType
_ SrcLoc
_) (UpdateSection [UpdateStep Info VName]
steps2 Info StructType
_ SrcLoc
_)
  | [UpdateStep Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [UpdateStep Info VName]
steps1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [UpdateStep Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [UpdateStep Info VName]
steps2 = [[(Exp, Exp)]] -> [(Exp, Exp)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Exp, Exp)]] -> [(Exp, Exp)])
-> Maybe [[(Exp, Exp)]] -> Maybe [(Exp, Exp)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (UpdateStep Info VName
 -> UpdateStep Info VName -> Maybe [(Exp, Exp)])
-> [UpdateStep Info VName]
-> [UpdateStep Info VName]
-> Maybe [[(Exp, Exp)]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM UpdateStep Info VName
-> UpdateStep Info VName -> Maybe [(Exp, Exp)]
similarStep [UpdateStep Info VName]
steps1 [UpdateStep Info VName]
steps2
  where
    similarStep :: UpdateStep Info VName
-> UpdateStep Info VName -> Maybe [(Exp, Exp)]
similarStep (UpdateStepField Name
f1) (UpdateStepField Name
f2)
      | Name
f1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
f2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
    similarStep (UpdateStepSlice Slice
s1) (UpdateStepSlice Slice
s2) =
      Slice -> Slice -> Maybe [(Exp, Exp)]
similarSlices Slice
s1 Slice
s2
    similarStep UpdateStep Info VName
_ UpdateStep Info VName
_ = Maybe [(Exp, Exp)]
forall a. Maybe a
Nothing
similarExps Exp
_ Exp
_ = Maybe [(Exp, Exp)]
forall a. Maybe a
Nothing

-- | Are these the same expression as per recursively invoking
-- 'similarExps'?
sameExp :: Exp -> Exp -> Bool
sameExp :: Exp -> Exp -> Bool
sameExp Exp
e1 Exp
e2
  | Just [(Exp, Exp)]
es <- Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps Exp
e1 Exp
e2 =
      ((Exp, Exp) -> Bool) -> [(Exp, Exp)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Exp -> Exp -> Bool) -> (Exp, Exp) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Exp -> Exp -> Bool
sameExp) [(Exp, Exp)]
es
  | Bool
otherwise = Bool
False

-- | An identifier with type- and aliasing information.
type Ident = IdentBase Info VName

-- | An index with type information.
type DimIndex = DimIndexBase Info VName

-- | A slice with type information.
type Slice = SliceBase Info VName

-- | An expression with type information.
type Exp = ExpBase Info VName

-- | An application expression with type information.
type AppExp = AppExpBase Info VName

-- | A pattern with type information.
type Pat = PatBase Info VName

-- | An constant declaration with type information.
type ValBind = ValBindBase Info VName

-- | A type binding with type information.
type TypeBind = TypeBindBase Info VName

-- | A type-checked module binding.
type ModBind = ModBindBase Info VName

-- | A type-checked module type binding.
type ModTypeBind = ModTypeBindBase Info VName

-- | A type-checked module expression.
type ModExp = ModExpBase Info VName

-- | A type-checked module parameter.
type ModParam = ModParamBase Info VName

-- | A type-checked module type expression.
type ModTypeExp = ModTypeExpBase Info VName

-- | A type-checked declaration.
type Dec = DecBase Info VName

-- | A type-checked specification.
type Spec = SpecBase Info VName

-- | An Futhark program with type information.
type Prog = ProgBase Info VName

-- | A known type arg with shape annotations.
type StructTypeArg = TypeArg Size

-- | A type-checked type parameter.
type TypeParam = TypeParamBase VName

-- | A known scalar type with no shape annotations.
type ScalarType = ScalarTypeBase ()

-- | A type-checked case (of a match expression).
type Case = CaseBase Info VName

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

-- | An unchecked type expression.
type UncheckedTypeExp = TypeExp UncheckedExp 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 UncheckedModTypeExp = ModTypeExpBase 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 type binding with no type annotations.
type UncheckedTypeBind = TypeBindBase NoInfo Name

-- | A module type binding with no type annotations.
type UncheckedModTypeBind = ModTypeBindBase NoInfo Name

-- | A module binding with no type annotations.
type UncheckedModBind = ModBindBase 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