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

    -- * 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
    uniqueness,
    unique,
    diet,
    arrayRank,
    arrayShape,
    orderZero,
    unfoldFunType,
    foldFunType,
    typeVars,
    isAccType,

    -- * Operations on types
    peelArray,
    stripArray,
    arrayOf,
    arrayOfWithAliases,
    toStructural,
    toStruct,
    toRes,
    toParam,
    resToParam,
    paramToRes,
    toResRet,
    setUniqueness,
    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.Loc (Loc (..), posFile)
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
import System.FilePath (takeDirectory)

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

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

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

-- | Change the shape of a type to be just the rank.
noSizes :: TypeBase Size as -> TypeBase () as
noSizes :: forall as. TypeBase Exp as -> TypeBase () as
noSizes = (Exp -> ()) -> TypeBase Exp as -> TypeBase () as
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 as -> TypeBase () as)
-> (Exp -> ()) -> TypeBase Exp as -> TypeBase () as
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 -> FilePath
(Int -> DimPos -> ShowS)
-> (DimPos -> FilePath) -> ([DimPos] -> ShowS) -> Show DimPos
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DimPos -> ShowS
showsPrec :: Int -> DimPos -> ShowS
$cshow :: DimPos -> FilePath
show :: DimPos -> FilePath
$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 als.
  (Applicative f) =>
  (S.Set VName -> DimPos -> fdim -> f tdim) ->
  TypeBase fdim als ->
  f (TypeBase tdim als)
traverseDims :: forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName -> DimPos -> fdim -> f tdim
f = Set VName -> DimPos -> TypeBase fdim als -> f (TypeBase tdim als)
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
forall a. Monoid a => a
mempty DimPos
PosImmediate
  where
    go ::
      forall als'.
      S.Set VName ->
      DimPos ->
      TypeBase fdim als' ->
      f (TypeBase tdim als')
    go :: forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b t :: TypeBase fdim als'
t@Array {} =
      (fdim -> f tdim)
-> (als' -> f als') -> TypeBase fdim als' -> f (TypeBase tdim als')
forall (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) als' -> f als'
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase fdim als'
t
    go Set VName
bound DimPos
b (Scalar (Record Map Name (TypeBase fdim als')
fields)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> (Map Name (TypeBase tdim als') -> ScalarTypeBase tdim als')
-> Map Name (TypeBase tdim als')
-> TypeBase tdim als'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase tdim als') -> ScalarTypeBase tdim als'
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase tdim als') -> TypeBase tdim als')
-> f (Map Name (TypeBase tdim als')) -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase fdim als' -> f (TypeBase tdim als'))
-> Map Name (TypeBase fdim als')
-> f (Map Name (TypeBase tdim als'))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b) Map Name (TypeBase fdim als')
fields
    go Set VName
bound DimPos
b (Scalar (TypeVar als'
as QualName VName
tn [TypeArg fdim]
targs)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> f (ScalarTypeBase tdim als') -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (als'
-> QualName VName -> [TypeArg tdim] -> ScalarTypeBase tdim als'
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar als'
as QualName VName
tn ([TypeArg tdim] -> ScalarTypeBase tdim als')
-> f [TypeArg tdim] -> f (ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg fdim -> f (TypeArg tdim))
-> [TypeArg fdim] -> f [TypeArg tdim]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
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 als']
cs)) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> (Map Name [TypeBase tdim als'] -> ScalarTypeBase tdim als')
-> Map Name [TypeBase tdim als']
-> TypeBase tdim als'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [TypeBase tdim als'] -> ScalarTypeBase tdim als'
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase tdim als'] -> TypeBase tdim als')
-> f (Map Name [TypeBase tdim als']) -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase fdim als'] -> f [TypeBase tdim als'])
-> Map Name [TypeBase fdim als']
-> f (Map Name [TypeBase tdim als'])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse ((TypeBase fdim als' -> f (TypeBase tdim als'))
-> [TypeBase fdim als'] -> f [TypeBase tdim als']
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b)) Map Name [TypeBase fdim als']
cs
    go Set VName
_ DimPos
_ (Scalar (Prim PrimType
t)) =
      TypeBase tdim als' -> f (TypeBase tdim als')
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase tdim als' -> f (TypeBase tdim als'))
-> TypeBase tdim als' -> f (TypeBase tdim als')
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> ScalarTypeBase tdim als' -> TypeBase tdim als'
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase tdim als'
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t
    go Set VName
bound DimPos
_ (Scalar (Arrow als'
als PName
p Diet
u TypeBase fdim NoUniqueness
t1 (RetType [VName]
dims TypeBase fdim Uniqueness
t2))) =
      ScalarTypeBase tdim als' -> TypeBase tdim als'
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase tdim als' -> TypeBase tdim als')
-> f (ScalarTypeBase tdim als') -> f (TypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (als'
-> PName
-> Diet
-> TypeBase tdim NoUniqueness
-> RetTypeBase tdim Uniqueness
-> ScalarTypeBase tdim als'
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow als'
als PName
p Diet
u (TypeBase tdim NoUniqueness
 -> RetTypeBase tdim Uniqueness -> ScalarTypeBase tdim als')
-> f (TypeBase tdim NoUniqueness)
-> f (RetTypeBase tdim Uniqueness -> ScalarTypeBase tdim als')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> DimPos
-> TypeBase fdim NoUniqueness
-> f (TypeBase tdim NoUniqueness)
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosParam TypeBase fdim NoUniqueness
t1 f (RetTypeBase tdim Uniqueness -> ScalarTypeBase tdim als')
-> f (RetTypeBase tdim Uniqueness) -> f (ScalarTypeBase tdim als')
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 Uniqueness -> RetTypeBase tdim Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase tdim Uniqueness -> RetTypeBase tdim Uniqueness)
-> f (TypeBase tdim Uniqueness) -> f (RetTypeBase tdim Uniqueness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> DimPos
-> TypeBase fdim Uniqueness
-> f (TypeBase tdim Uniqueness)
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound' DimPos
PosReturn TypeBase fdim Uniqueness
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 NoUniqueness
t) =
      TypeBase tdim NoUniqueness -> TypeArg tdim
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (TypeBase tdim NoUniqueness -> TypeArg tdim)
-> f (TypeBase tdim NoUniqueness) -> f (TypeArg tdim)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> DimPos
-> TypeBase fdim NoUniqueness
-> f (TypeBase tdim NoUniqueness)
forall als'.
Set VName -> DimPos -> TypeBase fdim als' -> f (TypeBase tdim als')
go Set VName
bound DimPos
b' TypeBase fdim NoUniqueness
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 uniqueness of a type.
uniqueness :: TypeBase shape Uniqueness -> Uniqueness
uniqueness :: forall shape. TypeBase shape Uniqueness -> Uniqueness
uniqueness (Array Uniqueness
u Shape shape
_ ScalarTypeBase shape NoUniqueness
_) = Uniqueness
u
uniqueness (Scalar (TypeVar Uniqueness
u QualName VName
_ [TypeArg shape]
_)) = Uniqueness
u
uniqueness (Scalar (Sum Map Name [TypeBase shape Uniqueness]
ts))
  | ([TypeBase shape Uniqueness] -> Bool)
-> Map Name [TypeBase shape Uniqueness] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((TypeBase shape Uniqueness -> Bool)
-> [TypeBase shape Uniqueness] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeBase shape Uniqueness -> Bool
forall shape. TypeBase shape Uniqueness -> Bool
unique) Map Name [TypeBase shape Uniqueness]
ts = Uniqueness
Unique
uniqueness (Scalar (Record Map Name (TypeBase shape Uniqueness)
fs))
  | (TypeBase shape Uniqueness -> Bool)
-> Map Name (TypeBase shape Uniqueness) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeBase shape Uniqueness -> Bool
forall shape. TypeBase shape Uniqueness -> Bool
unique Map Name (TypeBase shape Uniqueness)
fs = Uniqueness
Unique
uniqueness TypeBase shape Uniqueness
_ = Uniqueness
Nonunique

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

-- | @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 NoUniqueness
_) = 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 as ->
  TypeBase () ()
toStructural :: forall dim as. TypeBase dim as -> TypeBase () ()
toStructural = (dim -> ()) -> (as -> ()) -> TypeBase dim as -> TypeBase () ()
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 ()) (() -> as -> ()
forall a b. a -> b -> a
const ())

-- | Remove uniquenss information from a type.
toStruct :: TypeBase dim u -> TypeBase dim NoUniqueness
toStruct :: forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct = (u -> NoUniqueness) -> TypeBase dim u -> TypeBase dim NoUniqueness
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 (NoUniqueness -> u -> NoUniqueness
forall a b. a -> b -> a
const NoUniqueness
NoUniqueness)

-- | Uses 'Observe'.
toParam :: Diet -> TypeBase Size u -> ParamType
toParam :: forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
d = (u -> Diet) -> TypeBase Exp u -> 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 -> u -> Diet
forall a b. a -> b -> a
const Diet
d)

-- | Convert to 'ResType'
toRes :: Uniqueness -> TypeBase Size u -> ResType
toRes :: forall u. Uniqueness -> TypeBase Exp u -> ResType
toRes Uniqueness
u = (u -> Uniqueness) -> TypeBase Exp u -> 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 (Uniqueness -> u -> Uniqueness
forall a b. a -> b -> a
const Uniqueness
u)

-- | Convert to 'ResRetType'
toResRet :: Uniqueness -> RetTypeBase Size u -> ResRetType
toResRet :: forall u. Uniqueness -> RetTypeBase Exp u -> ResRetType
toResRet Uniqueness
u = (u -> Uniqueness) -> RetTypeBase Exp u -> 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 (Uniqueness -> u -> Uniqueness
forall a b. a -> b -> a
const Uniqueness
u)

-- | Preserves relation between 'Diet' and 'Uniqueness'.
resToParam :: ResType -> ParamType
resToParam :: ResType -> ParamType
resToParam = (Uniqueness -> 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 Uniqueness -> Diet
f
  where
    f :: Uniqueness -> Diet
f Uniqueness
Unique = Diet
Consume
    f Uniqueness
Nonunique = Diet
Observe

-- | Preserves relation between 'Diet' and 'Uniqueness'.
paramToRes :: ParamType -> ResType
paramToRes :: ParamType -> ResType
paramToRes = (Diet -> Uniqueness) -> 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 -> Uniqueness
f
  where
    f :: Diet -> Uniqueness
f Diet
Consume = Uniqueness
Unique
    f Diet
Observe = Uniqueness
Nonunique

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

-- | @arrayOf u s t@ constructs an array type.  The convenience
-- compared to using the 'Array' constructor directly is that @t@ can
-- itself be an array.  If @t@ is an @n@-dimensional array, and @s@ is
-- a list of length @n@, the resulting type is of an @n+m@ dimensions.
arrayOf ::
  Shape dim ->
  TypeBase dim NoUniqueness ->
  TypeBase dim NoUniqueness
arrayOf :: forall dim.
Shape dim -> TypeBase dim NoUniqueness -> TypeBase dim NoUniqueness
arrayOf = NoUniqueness
-> Shape dim
-> TypeBase dim NoUniqueness
-> TypeBase dim NoUniqueness
forall u dim u'.
u -> Shape dim -> TypeBase dim u' -> TypeBase dim u
arrayOfWithAliases NoUniqueness
forall a. Monoid a => a
mempty

-- | Like 'arrayOf', but you can pass in uniqueness info of the
-- resulting array.
arrayOfWithAliases ::
  u ->
  Shape dim ->
  TypeBase dim u' ->
  TypeBase dim u
arrayOfWithAliases :: forall u dim u'.
u -> Shape dim -> TypeBase dim u' -> TypeBase dim u
arrayOfWithAliases u
u Shape dim
shape2 (Array u'
_ Shape dim
shape1 ScalarTypeBase dim NoUniqueness
et) =
  u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u (Shape dim
shape2 Shape dim -> Shape dim -> Shape dim
forall a. Semigroup a => a -> a -> a
<> Shape dim
shape1) ScalarTypeBase dim NoUniqueness
et
arrayOfWithAliases u
u Shape dim
shape (Scalar ScalarTypeBase dim u'
t) =
  u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u Shape dim
shape ((u' -> NoUniqueness)
-> ScalarTypeBase dim u' -> ScalarTypeBase dim NoUniqueness
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 (NoUniqueness -> u' -> NoUniqueness
forall a b. a -> b -> a
const NoUniqueness
forall a. Monoid a => a
mempty) ScalarTypeBase dim u'
t)

-- | @stripArray n t@ removes the @n@ outermost layers of the array.
-- Essentially, it is the type of indexing an array of type @t@ with
-- @n@ indexes.
stripArray :: Int -> TypeBase dim as -> TypeBase dim as
stripArray :: forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
n (Array as
u Shape dim
shape ScalarTypeBase dim NoUniqueness
et)
  | Just Shape dim
shape' <- Int -> Shape dim -> Maybe (Shape dim)
forall dim. Int -> Shape dim -> Maybe (Shape dim)
stripDims Int
n Shape dim
shape =
      as
-> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim as
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array as
u Shape dim
shape' ScalarTypeBase dim NoUniqueness
et
  | Bool
otherwise =
      (NoUniqueness -> as)
-> TypeBase dim NoUniqueness -> TypeBase dim as
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 (as -> NoUniqueness -> as
forall a b. a -> b -> a
const as
u) (ScalarTypeBase dim NoUniqueness -> TypeBase dim NoUniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase dim NoUniqueness
et)
stripArray Int
_ TypeBase dim as
t = TypeBase dim as
t

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

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

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

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

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

-- | 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.  When present in
-- a type, each instance represents a distinct size.  The type checker
-- should _never_ produce these - 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 :: Size
anySize :: Exp
anySize =
  -- 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 [Word8
65, Word8
78, Word8
89] SrcLoc
forall a. Monoid a => a
mempty

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

    matchTypeArg :: [VName] -> TypeArg d1 -> TypeArg d2 -> m (TypeArg d1)
matchTypeArg [VName]
bound (TypeArgType TypeBase d1 NoUniqueness
t1) (TypeArgType TypeBase d2 NoUniqueness
t2) =
      TypeBase d1 NoUniqueness -> TypeArg d1
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (TypeBase d1 NoUniqueness -> TypeArg d1)
-> m (TypeBase d1 NoUniqueness) -> m (TypeArg d1)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName]
-> TypeBase d1 NoUniqueness
-> TypeBase d2 NoUniqueness
-> m (TypeBase d1 NoUniqueness)
forall u'.
Monoid u' =>
[VName] -> TypeBase d1 u' -> TypeBase d2 u' -> m (TypeBase d1 u')
matchDims' [VName]
bound TypeBase d1 NoUniqueness
t1 TypeBase d2 NoUniqueness
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 uniqueness attribute of a type.  If the type is a record
-- or sum type, the uniqueness of its components will be modified.
setUniqueness :: TypeBase dim u1 -> u2 -> TypeBase dim u2
setUniqueness :: forall dim u1 u2. TypeBase dim u1 -> u2 -> TypeBase dim u2
setUniqueness TypeBase dim u1
t u2
u = (u1 -> u2) -> TypeBase dim u1 -> TypeBase dim u2
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 (u2 -> u1 -> u2
forall a b. a -> b -> a
const u2
u) TypeBase dim u1
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp NoUniqueness)
-> PrimType -> ScalarTypeBase Exp NoUniqueness
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ [StructType] -> ScalarTypeBase Exp NoUniqueness
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord ([StructType] -> ScalarTypeBase Exp NoUniqueness)
-> [StructType] -> ScalarTypeBase Exp NoUniqueness
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Exp NoUniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name StructType -> ScalarTypeBase Exp NoUniqueness)
-> Map Name StructType -> ScalarTypeBase Exp NoUniqueness
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) =
  NoUniqueness
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array NoUniqueness
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 NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t)
typeOf (StringLit [Word8]
vs SrcLoc
loc) =
  NoUniqueness
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array
    NoUniqueness
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 NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
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
e SliceBase Info VName
_ Exp
_ SrcLoc
_) = Exp -> StructType
typeOf Exp
e
typeOf (RecordUpdate Exp
_ [Name]
_ 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
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
pn (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
pt2) (ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
pn (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
pt1) (ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
pt1) ResRetType
ret
typeOf (ProjectSection [Name]
_ (Info StructType
t) SrcLoc
_) = StructType
t
typeOf (IndexSection SliceBase 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 ParamType] -> ResRetType -> StructType
funType :: [PatBase Info VName ParamType] -> ResRetType -> StructType
funType [PatBase Info VName ParamType]
params ResRetType
ret =
  let RetType [VName]
_ ResType
t = (PatBase Info VName ParamType -> ResRetType -> ResRetType)
-> ResRetType -> [PatBase Info VName ParamType] -> ResRetType
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, StructType) -> ResRetType -> ResRetType
forall {dim}.
(PName, Diet, TypeBase dim NoUniqueness)
-> RetTypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
arrow ((PName, Diet, StructType) -> ResRetType -> ResRetType)
-> (PatBase Info VName ParamType -> (PName, Diet, StructType))
-> PatBase Info VName ParamType
-> ResRetType
-> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName ParamType -> (PName, Diet, StructType)
patternParam) ResRetType
ret [PatBase Info VName ParamType]
params
   in ResType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ResType
t
  where
    arrow :: (PName, Diet, TypeBase dim NoUniqueness)
-> RetTypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
arrow (PName
xp, Diet
d, TypeBase dim NoUniqueness
xt) RetTypeBase dim Uniqueness
yt =
      [VName] -> TypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim Uniqueness -> RetTypeBase dim Uniqueness)
-> TypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness)
-> ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ Uniqueness
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim Uniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow Uniqueness
Nonunique PName
xp Diet
d TypeBase dim NoUniqueness
xt RetTypeBase dim Uniqueness
yt

-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes
-- @ts@ as parameters and returns @ret@.
foldFunType :: [ParamType] -> ResRetType -> StructType
foldFunType :: [ParamType] -> ResRetType -> StructType
foldFunType [ParamType]
ps ResRetType
ret =
  let RetType [VName]
_ ResType
t = (ParamType -> ResRetType -> ResRetType)
-> ResRetType -> [ParamType] -> ResRetType
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ParamType -> ResRetType -> ResRetType
forall {dim}.
TypeBase dim Diet
-> RetTypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
arrow ResRetType
ret [ParamType]
ps
   in ResType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ResType
t
  where
    arrow :: TypeBase dim Diet
-> RetTypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
arrow TypeBase dim Diet
t1 RetTypeBase dim Uniqueness
t2 =
      [VName] -> TypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase dim Uniqueness -> RetTypeBase dim Uniqueness)
-> TypeBase dim Uniqueness -> RetTypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness)
-> ScalarTypeBase dim Uniqueness -> TypeBase dim Uniqueness
forall a b. (a -> b) -> a -> b
$ Uniqueness
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim Uniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow Uniqueness
Nonunique PName
Unnamed (TypeBase dim Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet TypeBase dim Diet
t1) (TypeBase dim Diet -> TypeBase dim NoUniqueness
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase dim Diet
t1) RetTypeBase dim Uniqueness
t2

-- | Extract the parameter types and return type from a type.
-- If the type is not an arrow type, the list of parameter types is empty.
unfoldFunType :: TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType :: forall dim as.
TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType (Scalar (Arrow as
_ PName
_ Diet
d TypeBase dim NoUniqueness
t1 (RetType [VName]
_ TypeBase dim Uniqueness
t2))) =
  let ([TypeBase dim Diet]
ps, TypeBase dim NoUniqueness
r) = TypeBase dim Uniqueness
-> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
forall dim as.
TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType TypeBase dim Uniqueness
t2
   in ((NoUniqueness -> Diet)
-> TypeBase dim NoUniqueness -> 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 -> NoUniqueness -> Diet
forall a b. a -> b -> a
const Diet
d) TypeBase dim NoUniqueness
t1 TypeBase dim Diet -> [TypeBase dim Diet] -> [TypeBase dim Diet]
forall a. a -> [a] -> [a]
: [TypeBase dim Diet]
ps, TypeBase dim NoUniqueness
r)
unfoldFunType TypeBase dim as
t = ([], TypeBase dim as -> TypeBase dim NoUniqueness
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase dim as
t)

-- | The type scheme of a value binding, comprising the type
-- parameters and the actual type.
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBindBase Info VName
vb =
  ( ValBindBase Info VName -> [TypeParamBase VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [TypeParamBase vn]
valBindTypeParams ValBindBase Info VName
vb,
    [PatBase Info VName ParamType] -> ResRetType -> StructType
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 as. RetTypeBase dim as -> [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 type names mentioned in a type.
typeVars :: TypeBase dim as -> S.Set VName
typeVars :: forall dim as. TypeBase dim as -> Set VName
typeVars TypeBase dim as
t =
  case TypeBase dim as
t of
    Scalar Prim {} -> Set VName
forall a. Monoid a => a
mempty
    Scalar (TypeVar as
_ QualName VName
tn [TypeArg dim]
targs) ->
      [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ VName -> Set VName
forall a. a -> Set a
S.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
tn) Set VName -> [Set VName] -> [Set VName]
forall a. a -> [a] -> [a]
: (TypeArg dim -> Set VName) -> [TypeArg dim] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg dim -> Set VName
forall {dim}. TypeArg dim -> Set VName
typeArgFree [TypeArg dim]
targs
    Scalar (Arrow as
_ PName
_ Diet
_ TypeBase dim NoUniqueness
t1 (RetType [VName]
_ TypeBase dim Uniqueness
t2)) -> TypeBase dim NoUniqueness -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars TypeBase dim NoUniqueness
t1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> TypeBase dim Uniqueness -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars TypeBase dim Uniqueness
t2
    Scalar (Record Map Name (TypeBase dim as)
fields) -> (TypeBase dim as -> Set VName)
-> Map Name (TypeBase dim as) -> Set 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 as -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars Map Name (TypeBase dim as)
fields
    Scalar (Sum Map Name [TypeBase dim as]
cs) -> [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (([TypeBase dim as] -> [Set VName])
-> Map Name [TypeBase dim as] -> [Set VName]
forall 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 as] -> [Set VName])
 -> Map Name [TypeBase dim as] -> [Set VName])
-> ((TypeBase dim as -> Set VName)
    -> [TypeBase dim as] -> [Set VName])
-> (TypeBase dim as -> Set VName)
-> Map Name [TypeBase dim as]
-> [Set VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase dim as -> Set VName) -> [TypeBase dim as] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) TypeBase dim as -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars Map Name [TypeBase dim as]
cs
    Array as
_ Shape dim
_ ScalarTypeBase dim NoUniqueness
rt -> TypeBase dim NoUniqueness -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars (TypeBase dim NoUniqueness -> Set VName)
-> TypeBase dim NoUniqueness -> Set VName
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim NoUniqueness -> TypeBase dim NoUniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase dim NoUniqueness
rt
  where
    typeArgFree :: TypeArg dim -> Set VName
typeArgFree (TypeArgType TypeBase dim NoUniqueness
ta) = TypeBase dim NoUniqueness -> Set VName
forall dim as. TypeBase dim as -> Set VName
typeVars TypeBase dim NoUniqueness
ta
    typeArgFree TypeArgDim {} = Set VName
forall a. Monoid a => a
mempty

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

-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern
-- have order 0.
patternOrderZero :: Pat (TypeBase d u) -> Bool
patternOrderZero :: forall d u. Pat (TypeBase d u) -> Bool
patternOrderZero = TypeBase d u -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (TypeBase d u -> Bool)
-> (Pat (TypeBase d u) -> TypeBase d u)
-> Pat (TypeBase d u)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (TypeBase d u) -> TypeBase d u
forall d u. Pat (TypeBase d u) -> TypeBase d u
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 u) -> TypeBase d u
patternType :: forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType (Wildcard (Info TypeBase d u
t) SrcLoc
_) = TypeBase d u
t
patternType (PatParens PatBase Info VName (TypeBase d u)
p SrcLoc
_) = PatBase Info VName (TypeBase d u) -> TypeBase d u
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType PatBase Info VName (TypeBase d u)
p
patternType (Id VName
_ (Info TypeBase d u
t) SrcLoc
_) = TypeBase d u
t
patternType (TuplePat [PatBase Info VName (TypeBase d u)]
pats SrcLoc
_) = ScalarTypeBase d u -> TypeBase d u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase d u -> TypeBase d u)
-> ScalarTypeBase d u -> TypeBase d u
forall a b. (a -> b) -> a -> b
$ [TypeBase d u] -> ScalarTypeBase d u
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord ([TypeBase d u] -> ScalarTypeBase d u)
-> [TypeBase d u] -> ScalarTypeBase d u
forall a b. (a -> b) -> a -> b
$ (PatBase Info VName (TypeBase d u) -> TypeBase d u)
-> [PatBase Info VName (TypeBase d u)] -> [TypeBase d u]
forall a b. (a -> b) -> [a] -> [b]
map PatBase Info VName (TypeBase d u) -> TypeBase d u
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType [PatBase Info VName (TypeBase d u)]
pats
patternType (RecordPat [(L Name, PatBase Info VName (TypeBase d u))]
fs SrcLoc
_) =
  ScalarTypeBase d u -> TypeBase d u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase d u -> TypeBase d u)
-> ScalarTypeBase d u -> TypeBase d u
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase d u) -> ScalarTypeBase d u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase d u) -> ScalarTypeBase d u)
-> Map Name (TypeBase d u) -> ScalarTypeBase d u
forall a b. (a -> b) -> a -> b
$ PatBase Info VName (TypeBase d u) -> TypeBase d u
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType (PatBase Info VName (TypeBase d u) -> TypeBase d u)
-> Map Name (PatBase Info VName (TypeBase d u))
-> Map Name (TypeBase d u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Name, PatBase Info VName (TypeBase d u))]
-> Map Name (PatBase Info VName (TypeBase d u))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (((L Name, PatBase Info VName (TypeBase d u))
 -> (Name, PatBase Info VName (TypeBase d u)))
-> [(L Name, PatBase Info VName (TypeBase d u))]
-> [(Name, PatBase Info VName (TypeBase d u))]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, PatBase Info VName (TypeBase d u))
-> (Name, PatBase Info VName (TypeBase d u))
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 u))]
fs)
patternType (PatAscription PatBase Info VName (TypeBase d u)
p TypeExp Exp VName
_ SrcLoc
_) = PatBase Info VName (TypeBase d u) -> TypeBase d u
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType PatBase Info VName (TypeBase d u)
p
patternType (PatLit PatLit
_ (Info TypeBase d u
t) SrcLoc
_) = TypeBase d u
t
patternType (PatConstr Name
_ (Info TypeBase d u
t) [PatBase Info VName (TypeBase d u)]
_ SrcLoc
_) = TypeBase d u
t
patternType (PatAttr AttrInfo VName
_ PatBase Info VName (TypeBase d u)
p SrcLoc
_) = PatBase Info VName (TypeBase d u) -> TypeBase d u
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType PatBase Info VName (TypeBase d u)
p

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

-- | When viewed as a function parameter, does this pattern correspond
-- to a named parameter of some type?
patternParam :: Pat ParamType -> (PName, Diet, StructType)
patternParam :: PatBase Info VName ParamType -> (PName, Diet, StructType)
patternParam (PatParens PatBase Info VName ParamType
p SrcLoc
_) =
  PatBase Info VName ParamType -> (PName, Diet, StructType)
patternParam PatBase Info VName ParamType
p
patternParam (PatAttr AttrInfo VName
_ PatBase Info VName ParamType
p SrcLoc
_) =
  PatBase Info VName ParamType -> (PName, Diet, StructType)
patternParam PatBase Info VName ParamType
p
patternParam (PatAscription (Id VName
v (Info ParamType
t) SrcLoc
_) TypeExp Exp VName
_ SrcLoc
_) =
  (VName -> PName
Named VName
v, ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
t, ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
t)
patternParam (Id VName
v (Info ParamType
t) SrcLoc
_) =
  (VName -> PName
Named VName
v, ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
t, ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
t)
patternParam PatBase Info VName ParamType
p =
  (PName
Unnamed, ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
p_t, ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
p_t)
  where
    p_t :: ParamType
p_t = PatBase Info VName ParamType -> ParamType
forall d u. Pat (TypeBase d u) -> TypeBase d u
patternType PatBase Info VName ParamType
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
    [ (FilePath -> Name
nameFromString (FilePath -> Name) -> FilePath -> Name
forall a b. (a -> b) -> a -> b
$ PrimType -> FilePath
forall a. Pretty a => a -> FilePath
prettyString 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 Uniqueness)
  | 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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$
        NoUniqueness
-> QualName VName
-> [TypeArg Exp]
-> ScalarTypeBase Exp NoUniqueness
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar NoUniqueness
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 NoUniqueness -> TypeArg dim
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (TypeBase dim NoUniqueness -> TypeArg dim)
-> TypeBase dim NoUniqueness -> TypeArg dim
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim NoUniqueness -> TypeBase dim NoUniqueness
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim NoUniqueness
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar NoUniqueness
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 u -> Maybe (TypeBase d NoUniqueness)
isAccType :: forall d u. TypeBase d u -> Maybe (TypeBase d NoUniqueness)
isAccType (Scalar (TypeVar u
_ (QualName [] VName
v) [TypeArgType TypeBase d NoUniqueness
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 NoUniqueness -> Maybe (TypeBase d NoUniqueness)
forall a. a -> Maybe a
Just TypeBase d NoUniqueness
t
isAccType TypeBase d u
_ = Maybe (TypeBase d NoUniqueness)
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 {a}. a
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 -> 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 :: a
bad = FilePath -> a
forall a. HasCallStack => FilePath -> a
error (FilePath -> a) -> FilePath -> a
forall a b. (a -> b) -> a -> b
$ FilePath
"findBuiltin: " FilePath -> ShowS
forall a. Semigroup a => a -> a -> a
<> Name -> FilePath
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp NoUniqueness)
-> PrimType -> ScalarTypeBase Exp NoUniqueness
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp NoUniqueness)
-> PrimType -> ScalarTypeBase Exp NoUniqueness
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 -> 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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar
                  (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
forall a. Monoid a => a
mempty
              ),
              ( 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 NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n, VName
m]) (ScalarTypeBase Exp NoUniqueness -> ParamType)
-> ScalarTypeBase Exp NoUniqueness -> ParamType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ResType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array
                    Uniqueness
Nonunique
                    ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
Shape [VName -> Exp
size VName
n Exp -> Exp -> Exp
`mkMul` VName -> Exp
size VName
m])
                    (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
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 NoUniqueness -> ParamType)
-> ScalarTypeBase Exp NoUniqueness -> ParamType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ResType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Uniqueness
Nonunique ([VName] -> Shape Exp
shape [VName
n, VName
m]) (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Nonunique
                  (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 NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Consume ([VName] -> Shape Exp
shape [VName
n]) (ScalarTypeBase Exp NoUniqueness -> ParamType)
-> ScalarTypeBase Exp NoUniqueness -> ParamType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp NoUniqueness)
-> PrimType -> ScalarTypeBase Exp NoUniqueness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (ScalarTypeBase Exp NoUniqueness -> ParamType)
-> ScalarTypeBase Exp NoUniqueness -> ParamType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ResType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Uniqueness
Unique ([VName] -> Shape Exp
shape [VName
n]) (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (Int -> ScalarTypeBase Exp NoUniqueness
forall {dim} {u}. Int -> ScalarTypeBase dim u
tupInt64 Int
2),
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (ScalarTypeBase Exp NoUniqueness -> ParamType)
-> ScalarTypeBase Exp NoUniqueness -> ParamType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (Int -> ScalarTypeBase Exp NoUniqueness
forall {dim} {u}. Int -> ScalarTypeBase dim u
tupInt64 Int
3),
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
l]) (ScalarTypeBase Exp NoUniqueness -> ParamType)
-> ScalarTypeBase Exp NoUniqueness -> ParamType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Diet
Observe ([VName] -> Shape Exp
shape [VName
n]),
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> StructType -> Shape Exp -> ResType
forall {u} {dim}.
u
-> TypeBase dim NoUniqueness
-> TypeBase dim NoUniqueness
-> Shape dim
-> TypeBase dim u
tuple_array Uniqueness
Unique (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b NoUniqueness
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 {u} {dim}.
u
-> TypeBase dim NoUniqueness
-> TypeBase dim NoUniqueness
-> Shape dim
-> TypeBase dim u
tuple_array Diet
Observe (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b NoUniqueness
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (ResType -> ResRetType)
-> ([(Name, ResType)] -> ResType)
-> [(Name, ResType)]
-> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Uniqueness -> ResType)
-> ([(Name, ResType)] -> ScalarTypeBase Exp Uniqueness)
-> [(Name, ResType)]
-> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name ResType -> ScalarTypeBase Exp Uniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name ResType -> ScalarTypeBase Exp Uniqueness)
-> ([(Name, ResType)] -> Map Name ResType)
-> [(Name, ResType)]
-> ScalarTypeBase Exp Uniqueness
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 [Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n], Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_b Uniqueness
Unique (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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n]) (Int -> ScalarTypeBase Exp NoUniqueness
forall {dim} {u}. Int -> ScalarTypeBase dim u
tupInt64 Int
1),
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n]) (Int -> ScalarTypeBase Exp NoUniqueness
forall {dim} {u}. Int -> ScalarTypeBase dim u
tupInt64 Int
2),
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    Diet -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ParamType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Diet
Observe ([VName] -> Shape Exp
shape [VName
n]) (Int -> ScalarTypeBase Exp NoUniqueness
forall {dim} {u}. Int -> ScalarTypeBase dim u
tupInt64 Int
3),
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Uniqueness
Nonunique),
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_b Uniqueness
Unique
                  (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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Unique)
              ),
              ( Text
"reduce_comm",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Unique))
              ),
              ( Text
"scan",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
sp_n]
                  [ ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique (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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp Diet)
-> PrimType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int32),
                    ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Uniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp Uniqueness)
-> PrimType -> ScalarTypeBase Exp Uniqueness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64),
                    Diet -> Shape Exp -> ParamType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
k] (ResType -> ResRetType)
-> (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness
-> ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Uniqueness -> ResRetType)
-> ScalarTypeBase Exp Uniqueness -> ResRetType
forall a b. (a -> b) -> a -> b
$
                      [ResType] -> ScalarTypeBase Exp Uniqueness
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord
                        [ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique (Shape Exp -> ResType) -> Shape Exp -> ResType
forall a b. (a -> b) -> a -> b
$ [VName] -> Shape Exp
shape [VName
n],
                          Uniqueness
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> ResType
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array Uniqueness
Unique ([VName] -> Shape Exp
shape [VName
k]) (PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim (PrimType -> ScalarTypeBase Exp NoUniqueness)
-> PrimType -> ScalarTypeBase Exp NoUniqueness
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> StructType -> ScalarTypeBase Exp Diet
forall {u} {dim}.
u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType Diet
Consume (StructType -> ScalarTypeBase Exp Diet)
-> StructType -> ScalarTypeBase Exp Diet
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> StructType
forall {u}. u -> TypeBase Exp u
array_ka NoUniqueness
forall a. Monoid a => a
mempty,
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar
                  (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}.
u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType Uniqueness
Unique (NoUniqueness -> StructType
forall {u}. u -> TypeBase Exp u
array_ka NoUniqueness
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 {u}. u -> TypeBase Exp u
array_ka Diet
Consume,
                    ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> StructType -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}.
u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType NoUniqueness
forall a. Monoid a => a
mempty (NoUniqueness -> StructType
forall {u}. u -> TypeBase Exp u
array_ka NoUniqueness
forall a. Monoid a => a
mempty))
                      StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`carr` ( ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b NoUniqueness
forall a. Monoid a => a
mempty)
                                 StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> StructType -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}.
u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType Uniqueness
Nonunique (StructType -> ScalarTypeBase Exp Uniqueness)
-> StructType -> ScalarTypeBase Exp Uniqueness
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> Shape Exp -> StructType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a NoUniqueness
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> ResType
forall {u}. u -> TypeBase Exp u
array_ka Uniqueness
Unique
              ),
              ( 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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique)),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Diet -> ParamType)
-> ScalarTypeBase Exp Diet -> ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe,
                    ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> StructType -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}.
u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType NoUniqueness
forall a. Monoid a => a
mempty (StructType -> ScalarTypeBase Exp NoUniqueness)
-> StructType -> ScalarTypeBase Exp NoUniqueness
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> StructType
forall {u}. u -> TypeBase Exp u
array_ka NoUniqueness
forall a. Monoid a => a
mempty)
                      StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`carr` ( ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b NoUniqueness
forall a. Monoid a => a
mempty)
                                 StructType -> ResType -> ResType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> StructType -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}.
u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType Uniqueness
Nonunique (StructType -> ScalarTypeBase Exp Uniqueness)
-> StructType -> ScalarTypeBase Exp Uniqueness
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> Shape Exp -> StructType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a NoUniqueness
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                  (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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Uniqueness
Nonunique),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe)
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar
                  (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ [ResType] -> ScalarTypeBase Exp Uniqueness
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord [ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Uniqueness
Nonunique, ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Uniqueness
Nonunique]
              ),
              ( Text
"vjp2",
                [TypeParamBase VName] -> [ParamType] -> ResRetType -> Intrinsic
IntrinsicPolyFun
                  [TypeParamBase VName
tp_a, TypeParamBase VName
tp_b]
                  [ ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty) StructType -> ResType -> ParamType
forall {u} {dim}.
Monoid u =>
TypeBase dim NoUniqueness
-> TypeBase dim Uniqueness -> TypeBase dim u
`arr` ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Uniqueness
Nonunique),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Diet
Observe),
                    ScalarTypeBase Exp Diet -> ParamType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (Diet -> ScalarTypeBase Exp Diet
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Diet
Observe)
                  ]
                  (ResRetType -> Intrinsic) -> ResRetType -> Intrinsic
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                  (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar
                  (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ [ResType] -> ScalarTypeBase Exp Uniqueness
forall dim as. [TypeBase dim as] -> ScalarTypeBase dim as
tupleRecord [ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_b Uniqueness
Nonunique, ScalarTypeBase Exp Uniqueness -> ResType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp Uniqueness -> ResType)
-> ScalarTypeBase Exp Uniqueness -> ResType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> ScalarTypeBase Exp Uniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a Uniqueness
Nonunique]
              )
            ]
              [(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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m, VName
k]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Nonunique
                    (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                    (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m, VName
k, VName
l]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Nonunique
                    (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                    (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName
m, VName
k, VName
l, VName
p]
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Nonunique
                    (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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (PrimType -> ScalarTypeBase Exp Diet
forall dim u. PrimType -> ScalarTypeBase dim u
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 {u} {dim}. u -> Shape dim -> TypeBase dim u
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 as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []
                    (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Shape Exp -> ResType
forall {u} {dim}. u -> Shape dim -> TypeBase dim u
array_a Uniqueness
Unique
                    (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
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
"n", Text
"m", Text
"k", Text
"l", Text
"p", Text
"q"]) [Int
0 ..]

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

    [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 :: u
-> TypeBase dim NoUniqueness
-> TypeBase dim NoUniqueness
-> Shape dim
-> TypeBase dim u
tuple_array u
u TypeBase dim NoUniqueness
x TypeBase dim NoUniqueness
y Shape dim
s =
      u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u Shape dim
s (Map Name (TypeBase dim NoUniqueness)
-> ScalarTypeBase dim NoUniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record ([(Name, TypeBase dim NoUniqueness)]
-> Map Name (TypeBase dim NoUniqueness)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeBase dim NoUniqueness)]
 -> Map Name (TypeBase dim NoUniqueness))
-> [(Name, TypeBase dim NoUniqueness)]
-> Map Name (TypeBase dim NoUniqueness)
forall a b. (a -> b) -> a -> b
$ [Name]
-> [TypeBase dim NoUniqueness]
-> [(Name, TypeBase dim NoUniqueness)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [TypeBase dim NoUniqueness
x, TypeBase dim NoUniqueness
y]))

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

    array_ka :: u -> TypeBase Exp u
array_ka u
u = u -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u ([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 NoUniqueness -> TypeBase Exp u)
-> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u
forall a b. (a -> b) -> a -> b
$ NoUniqueness -> ScalarTypeBase Exp NoUniqueness
forall {u} {dim}. u -> ScalarTypeBase dim u
t_a NoUniqueness
forall a. Monoid a => a
mempty

    accType :: u -> TypeBase dim NoUniqueness -> ScalarTypeBase dim u
accType u
u TypeBase dim NoUniqueness
t =
      u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar u
u (VName -> QualName VName
forall v. v -> QualName v
qualName ((VName, Intrinsic) -> VName
forall a b. (a, b) -> a
fst (VName, Intrinsic)
intrinsicAcc)) [TypeBase dim NoUniqueness -> TypeArg dim
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType TypeBase dim NoUniqueness
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 NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
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
      Intrinsic
op' <- BinOp -> Maybe Intrinsic
intrinsicBinOp BinOp
op
      (Text, Intrinsic) -> Maybe (Text, Intrinsic)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BinOp -> Text
forall a. Pretty a => a -> Text
prettyText BinOp
op, Intrinsic
op')

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

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

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

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

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

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

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

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

-- | The modules imported by a Futhark program.
progImports :: ProgBase f vn -> [(String, Loc)]
progImports :: forall (f :: * -> *) vn. ProgBase f vn -> [(FilePath, Loc)]
progImports = (DecBase f vn -> [(FilePath, Loc)])
-> [DecBase f vn] -> [(FilePath, Loc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(FilePath, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(FilePath, Loc)]
decImports ([DecBase f vn] -> [(FilePath, Loc)])
-> (ProgBase f vn -> [DecBase f vn])
-> ProgBase f vn
-> [(FilePath, 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 -> [(FilePath, Loc)]
decImports (OpenDec ModExpBase f vn
x SrcLoc
_) = ModExpBase f vn -> [(FilePath, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, Loc)]
modExpImports ModExpBase f vn
x
decImports (ModDec ModBindBase f vn
md) = ModExpBase f vn -> [(FilePath, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, Loc)]
modExpImports (ModExpBase f vn -> [(FilePath, Loc)])
-> ModExpBase f vn -> [(FilePath, 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 -> [(FilePath, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(FilePath, Loc)]
decImports DecBase f vn
d
decImports (ImportDec FilePath
x f ImportName
_ SrcLoc
loc) = [(FilePath
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 -> [(FilePath, Loc)]
modExpImports ModVar {} = []
modExpImports (ModParens ModExpBase f vn
p SrcLoc
_) = ModExpBase f vn -> [(FilePath, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, Loc)]
modExpImports ModExpBase f vn
p
modExpImports (ModImport FilePath
f f ImportName
_ SrcLoc
loc) = [(FilePath
f, SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc)]
modExpImports (ModDecs [DecBase f vn]
ds SrcLoc
_) = (DecBase f vn -> [(FilePath, Loc)])
-> [DecBase f vn] -> [(FilePath, Loc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(FilePath, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(FilePath, 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 -> [(FilePath, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, Loc)]
modExpImports ModExpBase f vn
me
modExpImports (ModAscript ModExpBase f vn
me ModTypeExpBase f vn
_ f (Map VName VName)
_ SrcLoc
_) = ModExpBase f vn -> [(FilePath, Loc)]
forall (f :: * -> *) vn. ModExpBase f vn -> [(FilePath, 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 :: FilePath -> Maybe ((FilePath, FilePath, Maybe FilePath), FilePath)
identifierReference (Char
'`' : FilePath
s)
  | (FilePath
identifier, Char
'`' : Char
'@' : FilePath
s') <- (Char -> Bool) -> FilePath -> (FilePath, FilePath)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'`') FilePath
s,
    (FilePath
namespace, FilePath
s'') <- (Char -> Bool) -> FilePath -> (FilePath, FilePath)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isAlpha FilePath
s',
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ FilePath -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null FilePath
namespace =
      case FilePath
s'' of
        Char
'@' : Char
'"' : FilePath
s'''
          | (FilePath
file, Char
'"' : FilePath
s'''') <- (Char -> Bool) -> FilePath -> (FilePath, FilePath)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'"') FilePath
s''' ->
              ((FilePath, FilePath, Maybe FilePath), FilePath)
-> Maybe ((FilePath, FilePath, Maybe FilePath), FilePath)
forall a. a -> Maybe a
Just ((FilePath
identifier, FilePath
namespace, FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
file), FilePath
s'''')
        FilePath
_ -> ((FilePath, FilePath, Maybe FilePath), FilePath)
-> Maybe ((FilePath, FilePath, Maybe FilePath), FilePath)
forall a. a -> Maybe a
Just ((FilePath
identifier, FilePath
namespace, Maybe FilePath
forall a. Maybe a
Nothing), FilePath
s'')
identifierReference FilePath
_ = Maybe ((FilePath, FilePath, Maybe FilePath), FilePath)
forall a. Maybe a
Nothing

-- | Given an operator name, return the operator that determines its
-- syntactical properties.
leadingOperator :: Name -> BinOp
leadingOperator :: Name -> BinOp
leadingOperator Name
s =
  BinOp
-> ((FilePath, BinOp) -> BinOp) -> Maybe (FilePath, BinOp) -> BinOp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe BinOp
Backtick (FilePath, BinOp) -> BinOp
forall a b. (a, b) -> b
snd (Maybe (FilePath, BinOp) -> BinOp)
-> Maybe (FilePath, BinOp) -> BinOp
forall a b. (a -> b) -> a -> b
$
    ((FilePath, BinOp) -> Bool)
-> [(FilePath, BinOp)] -> Maybe (FilePath, BinOp)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((FilePath -> FilePath -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
s') (FilePath -> Bool)
-> ((FilePath, BinOp) -> FilePath) -> (FilePath, BinOp) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath, BinOp) -> FilePath
forall a b. (a, b) -> a
fst) ([(FilePath, BinOp)] -> Maybe (FilePath, BinOp))
-> [(FilePath, BinOp)] -> Maybe (FilePath, BinOp)
forall a b. (a -> b) -> a -> b
$
      ((FilePath, BinOp) -> Down Int)
-> [(FilePath, BinOp)] -> [(FilePath, BinOp)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int -> Down Int
forall a. a -> Down a
Down (Int -> Down Int)
-> ((FilePath, BinOp) -> Int) -> (FilePath, BinOp) -> Down Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FilePath -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (FilePath -> Int)
-> ((FilePath, BinOp) -> FilePath) -> (FilePath, BinOp) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath, BinOp) -> FilePath
forall a b. (a, b) -> a
fst) ([(FilePath, BinOp)] -> [(FilePath, BinOp)])
-> [(FilePath, BinOp)] -> [(FilePath, BinOp)]
forall a b. (a -> b) -> a -> b
$
        [FilePath] -> [BinOp] -> [(FilePath, BinOp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((BinOp -> FilePath) -> [BinOp] -> [FilePath]
forall a b. (a -> b) -> [a] -> [b]
map BinOp -> FilePath
forall a. Pretty a => a -> FilePath
prettyString [BinOp]
operators) [BinOp]
operators
  where
    s' :: FilePath
s' = Name -> FilePath
nameToString Name
s
    operators :: [BinOp]
    operators :: [BinOp]
operators = [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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t) :)
      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' :)
          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 :: SliceBase Info VName -> SliceBase Info VName -> Maybe [(Exp, Exp)]
similarSlices SliceBase Info VName
slice1 SliceBase Info VName
slice2
  | SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
slice1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
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)])
-> SliceBase Info VName
-> SliceBase Info VName
-> 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 SliceBase Info VName
slice1 SliceBase Info VName
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 SliceBase Info VName
slice1 SrcLoc
_) Info AppRes
_) (AppExp (Index Exp
arr2 SliceBase Info VName
slice2 SrcLoc
_) Info AppRes
_)
  | Exp
arr1 Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
arr2,
    SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
slice1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
slice2 =
      SliceBase Info VName -> SliceBase Info VName -> Maybe [(Exp, Exp)]
similarSlices SliceBase Info VName
slice1 SliceBase Info VName
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 SliceBase Info VName
slice1 Exp
e'1 SrcLoc
_) (Update Exp
e2 SliceBase Info VName
slice2 Exp
e'2 SrcLoc
_) =
  ([(Exp
e1, Exp
e2), (Exp
e'1, Exp
e'2)] ++) ([(Exp, Exp)] -> [(Exp, Exp)])
-> Maybe [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SliceBase Info VName -> SliceBase Info VName -> Maybe [(Exp, Exp)]
similarSlices SliceBase Info VName
slice1 SliceBase Info VName
slice2
similarExps (RecordUpdate Exp
e1 [Name]
names1 Exp
e'1 Info StructType
_ SrcLoc
_) (RecordUpdate Exp
e2 [Name]
names2 Exp
e'2 Info StructType
_ SrcLoc
_)
  | [Name]
names1 [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name]
names2 =
      [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just [(Exp
e1, Exp
e2), (Exp
e'1, Exp
e'2)]
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 (ProjectSection [Name]
names1 Info StructType
_ SrcLoc
_) (ProjectSection [Name]
names2 Info StructType
_ SrcLoc
_)
  | [Name]
names1 [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== [Name]
names2 = [(Exp, Exp)] -> Maybe [(Exp, Exp)]
forall a. a -> Maybe a
Just []
similarExps (IndexSection SliceBase Info VName
slice1 Info StructType
_ SrcLoc
_) (IndexSection SliceBase Info VName
slice2 Info StructType
_ SrcLoc
_) =
  SliceBase Info VName -> SliceBase Info VName -> Maybe [(Exp, Exp)]
similarSlices SliceBase Info VName
slice1 SliceBase Info VName
slice2
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