-- | This module provides various simple ways to query and manipulate
-- fundamental Futhark terms, such as types and values.  The intent is to
-- keep "Futhark.Language.Syntax" simple, and put whatever embellishments
-- we need here.
module Language.Futhark.Prop
  ( -- * Various
    Intrinsic (..),
    intrinsics,
    intrinsicVar,
    maxIntrinsicTag,
    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.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.Util (maxinum)
import Futhark.Util.Pretty
import Language.Futhark.Primitive qualified as Primitive
import Language.Futhark.Syntax
import Language.Futhark.Traversals
import Language.Futhark.Tuple

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

-- | Return the dimensionality of a type.  For non-arrays, this is
-- zero.  For a one-dimensional array it is one, for a two-dimensional
-- it is two, and so forth.
arrayRank :: TypeBase d 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 -> String
(Int -> DimPos -> ShowS)
-> (DimPos -> String) -> ([DimPos] -> ShowS) -> Show DimPos
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DimPos -> ShowS
showsPrec :: Int -> DimPos -> ShowS
$cshow :: DimPos -> String
show :: DimPos -> String
$cshowList :: [DimPos] -> ShowS
showList :: [DimPos] -> ShowS
Show)

-- | Perform a traversal (possibly including replacement) on sizes
-- that are parameters in a function type, but also including the type
-- immediately passed to the function.  Also passes along a set of the
-- parameter names inside the type that have come in scope at the
-- occurrence of the dimension.
traverseDims ::
  forall f fdim tdim 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
    [ (String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ PrimType -> String
forall a. Pretty a => a -> String
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 = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"findBuiltin: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Name -> String
nameToString Name
v

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

mkAdd, mkMul :: Exp -> Exp -> Exp
mkAdd :: Exp -> Exp -> Exp
mkAdd = Name -> StructType -> Exp -> Exp -> Exp
mkBinOp Name
"+" (StructType -> Exp -> Exp -> Exp)
-> StructType -> Exp -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp 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

-- | 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 -> [(String, Loc)]
progImports = (DecBase f vn -> [(String, Loc)])
-> [DecBase f vn] -> [(String, Loc)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DecBase f vn -> [(String, Loc)]
forall (f :: * -> *) vn. DecBase f vn -> [(String, Loc)]
decImports ([DecBase f vn] -> [(String, Loc)])
-> (ProgBase f vn -> [DecBase f vn])
-> ProgBase f vn
-> [(String, Loc)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgBase f vn -> [DecBase f vn]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs

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

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

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

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

        onModTypeExp :: ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeVar QualName a
v f (Map VName VName)
_ SrcLoc
_) = a -> Set a
forall a. a -> Set a
S.singleton (a -> Set a) -> a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
v
        onModTypeExp (ModTypeParens ModTypeExpBase f a
e SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp (ModTypeSpecs [SpecBase f a]
ss SrcLoc
_) = (SpecBase f a -> Set a) -> [SpecBase f a] -> Set a
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap SpecBase f a -> Set a
onSpec [SpecBase f a]
ss
        onModTypeExp (ModTypeWith ModTypeExpBase f a
e TypeRefBase f a
_ SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp (ModTypeArrow Maybe a
_ ModTypeExpBase f a
e1 ModTypeExpBase f a
e2 SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e2

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

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

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

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

        onModTypeExp :: ModTypeExpBase f a -> Set a
onModTypeExp (ModTypeVar QualName a
v f (Map VName VName)
_ SrcLoc
_) = a -> Set a
forall a. a -> Set a
S.singleton (a -> Set a) -> a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
v
        onModTypeExp (ModTypeParens ModTypeExpBase f a
e SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp ModTypeSpecs {} = Set a
forall a. Monoid a => a
mempty
        onModTypeExp (ModTypeWith ModTypeExpBase f a
e TypeRefBase f a
_ SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e
        onModTypeExp (ModTypeArrow Maybe a
_ ModTypeExpBase f a
e1 ModTypeExpBase f a
e2 SrcLoc
_) = ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e1 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> ModTypeExpBase f a -> Set a
onModTypeExp ModTypeExpBase f a
e2

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

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

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

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

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

    onExp :: Exp -> m Exp
onExp e :: Exp
e@(Hole (Info StructType
t) SrcLoc
loc) = do
      ([(Loc, StructType)] -> [(Loc, StructType)]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc, StructType -> StructType
forall dim 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