-- | Functions for inspecting and constructing various types.
module Futhark.IR.Prop.Types
  ( rankShaped,
    arrayRank,
    arrayShape,
    setArrayShape,
    isEmptyArray,
    existential,
    consuming,
    staticShapes,
    staticShapes1,
    primType,
    isAcc,
    isMem,
    arrayOf,
    arrayOfRow,
    arrayOfShape,
    setOuterSize,
    setDimSize,
    setOuterDim,
    setOuterDims,
    setDim,
    setArrayDims,
    peelArray,
    stripArray,
    arrayDims,
    shapeSize,
    arraySize,
    arraysSize,
    elemType,
    rowType,
    transposeType,
    rearrangeType,
    mapOnExtType,
    mapOnType,
    diet,
    toDecl,
    fromDecl,
    isExt,
    isFree,
    extractShapeContext,
    shapeContext,
    hasStaticShape,
    generaliseExtTypes,
    existentialiseExtTypes,
    shapeExtMapping,

    -- * Abbreviations
    int8,
    int16,
    int32,
    int64,
    float32,
    float64,

    -- * The Typed typeclass
    Typed (..),
    DeclTyped (..),
    ExtTyped (..),
    FixExt (..),
  )
where

import Control.Monad
import Control.Monad.State
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.IR.Prop.Constants
import Futhark.IR.Prop.Rearrange
import Futhark.IR.Syntax.Core

-- | Remove shape information from a type.
rankShaped :: (ArrayShape shape) => TypeBase shape o -> TypeBase Rank o
rankShaped :: forall shape o.
ArrayShape shape =>
TypeBase shape o -> TypeBase Rank o
rankShaped (Array PrimType
et shape
sz o
o) = PrimType -> Rank -> o -> TypeBase Rank o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et (Int -> Rank
Rank (Int -> Rank) -> Int -> Rank
forall a b. (a -> b) -> a -> b
$ shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank shape
sz) o
o
rankShaped (Prim PrimType
pt) = PrimType -> TypeBase Rank o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt
rankShaped (Acc VName
acc Shape
ispace [Type]
ts) = VName -> Shape -> [Type] -> TypeBase Rank o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
rankShaped (Mem Space
space) = Space -> TypeBase Rank o
forall shape o. Space -> TypeBase shape o
Mem Space
space

-- | 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 :: (ArrayShape shape) => TypeBase shape o -> Int
arrayRank :: forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank = shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (shape -> Int)
-> (TypeBase shape o -> shape) -> TypeBase shape o -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase shape o -> shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape

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

-- | Modify the shape of an array - for non-arrays, this does nothing.
modifyArrayShape ::
  (ArrayShape newshape) =>
  (oldshape -> newshape) ->
  TypeBase oldshape o ->
  TypeBase newshape o
modifyArrayShape :: forall newshape oldshape o.
ArrayShape newshape =>
(oldshape -> newshape)
-> TypeBase oldshape o -> TypeBase newshape o
modifyArrayShape oldshape -> newshape
f (Array PrimType
t oldshape
ds o
o)
  | newshape -> Int
forall a. ArrayShape a => a -> Int
shapeRank newshape
ds' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = PrimType -> TypeBase newshape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t
  | Bool
otherwise = PrimType -> newshape -> o -> TypeBase newshape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
t newshape
ds' o
o
  where
    ds' :: newshape
ds' = oldshape -> newshape
f oldshape
ds
modifyArrayShape oldshape -> newshape
_ (Prim PrimType
t) = PrimType -> TypeBase newshape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t
modifyArrayShape oldshape -> newshape
_ (Acc VName
acc Shape
ispace [Type]
ts) = VName -> Shape -> [Type] -> TypeBase newshape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
modifyArrayShape oldshape -> newshape
_ (Mem Space
space) = Space -> TypeBase newshape o
forall shape o. Space -> TypeBase shape o
Mem Space
space

-- | Set the shape of an array.  If the given type is not an
-- array, return the type unchanged.
setArrayShape ::
  (ArrayShape newshape) =>
  TypeBase oldshape o ->
  newshape ->
  TypeBase newshape o
setArrayShape :: forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
setArrayShape TypeBase oldshape o
t newshape
ds = (oldshape -> newshape)
-> TypeBase oldshape o -> TypeBase newshape o
forall newshape oldshape o.
ArrayShape newshape =>
(oldshape -> newshape)
-> TypeBase oldshape o -> TypeBase newshape o
modifyArrayShape (newshape -> oldshape -> newshape
forall a b. a -> b -> a
const newshape
ds) TypeBase oldshape o
t

-- | If the array is statically an empty array (meaning any dimension
-- is a static zero), return the element type and the shape.
isEmptyArray :: Type -> Maybe (PrimType, Shape)
isEmptyArray :: Type -> Maybe (PrimType, Shape)
isEmptyArray (Array PrimType
pt (Shape [SubExp]
ds) NoMode
_)
  | IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0 SubExp -> [SubExp] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [SubExp]
ds = (PrimType, Shape) -> Maybe (PrimType, Shape)
forall a. a -> Maybe a
Just (PrimType
pt, [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
ds)
isEmptyArray Type
_ = Maybe (PrimType, Shape)
forall a. Maybe a
Nothing

-- | True if the given type has a dimension that is existentially sized.
existential :: ExtType -> Bool
existential :: ExtType -> Bool
existential = (Ext SubExp -> Bool) -> [Ext SubExp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Ext SubExp -> Bool
forall {a}. Ext a -> Bool
ext ([Ext SubExp] -> Bool)
-> (ExtType -> [Ext SubExp]) -> ExtType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtShape -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ExtShape -> [Ext SubExp])
-> (ExtType -> ExtShape) -> ExtType -> [Ext SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtType -> ExtShape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape
  where
    ext :: Ext a -> Bool
ext (Ext Int
_) = Bool
True
    ext (Free a
_) = Bool
False

-- | @consuming t@ is 'True' if a parameter of this type consumes its
-- argument.
consuming :: TypeBase shape Diet -> Bool
consuming :: forall shape. TypeBase shape Diet -> Bool
consuming = (Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
Consume) (Diet -> Bool)
-> (TypeBase shape Diet -> Diet) -> TypeBase shape Diet -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase shape Diet -> Diet
forall shape. TypeBase shape Diet -> Diet
diet

-- | Convert types with non-existential shapes to types with
-- existential shapes.  Only the representation is changed, so all
-- the shapes will be 'Free'.
staticShapes :: [TypeBase Shape o] -> [TypeBase ExtShape o]
staticShapes :: forall o. [TypeBase Shape o] -> [TypeBase ExtShape o]
staticShapes = (TypeBase Shape o -> TypeBase ExtShape o)
-> [TypeBase Shape o] -> [TypeBase ExtShape o]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape o -> TypeBase ExtShape o
forall o. TypeBase Shape o -> TypeBase ExtShape o
staticShapes1

-- | As 'staticShapes', but on a single type.
staticShapes1 :: TypeBase Shape o -> TypeBase ExtShape o
staticShapes1 :: forall o. TypeBase Shape o -> TypeBase ExtShape o
staticShapes1 (Prim PrimType
t) =
  PrimType -> TypeBase ExtShape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t
staticShapes1 (Acc VName
acc Shape
ispace [Type]
ts) =
  VName -> Shape -> [Type] -> TypeBase ExtShape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
staticShapes1 (Array PrimType
bt (Shape [SubExp]
shape) o
o) =
  PrimType -> ExtShape -> o -> TypeBase ExtShape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
bt ([Ext SubExp] -> ExtShape
forall d. [d] -> ShapeBase d
Shape ([Ext SubExp] -> ExtShape) -> [Ext SubExp] -> ExtShape
forall a b. (a -> b) -> a -> b
$ (SubExp -> Ext SubExp) -> [SubExp] -> [Ext SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Ext SubExp
forall a. a -> Ext a
Free [SubExp]
shape) o
o
staticShapes1 (Mem Space
space) =
  Space -> TypeBase ExtShape o
forall shape o. Space -> TypeBase shape o
Mem Space
space

-- | @arrayOf t s o@ constructs an array type.  The convenience
-- compared to using the 'Array' constructor directly is that @t@ can
-- itself be an array.  If @t@ is an @n@-dimensional array, and @s@ is
-- a list of length @n@, the resulting type is of an @n+m@ dimensions.
-- The mode of the new array will be @o@, no matter the mode of @t@.
-- If the shape @s@ has rank 0, then the @t@ will be returned,
-- although if it is an array, with the mode changed to @o@.
arrayOf ::
  (ArrayShape shape) =>
  TypeBase shape o_unused ->
  shape ->
  o ->
  TypeBase shape o
arrayOf :: forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf (Array PrimType
et shape
size1 o_unused
_) shape
size2 o
o =
  PrimType -> shape -> o -> TypeBase shape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et (shape
size2 shape -> shape -> shape
forall a. Semigroup a => a -> a -> a
<> shape
size1) o
o
arrayOf (Prim PrimType
t) shape
shape o
o
  | Int
0 <- shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank shape
shape = PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t
  | Bool
otherwise = PrimType -> shape -> o -> TypeBase shape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
t shape
shape o
o
arrayOf (Acc VName
acc Shape
ispace [Type]
ts) shape
_shape o
_ =
  VName -> Shape -> [Type] -> TypeBase shape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
arrayOf Mem {} shape
_ o
_ =
  [Char] -> TypeBase shape o
forall a. HasCallStack => [Char] -> a
error [Char]
"arrayOf Mem"

-- | Construct an array whose rows are the given type, and the outer
-- size is the given dimension.  This is just a convenient wrapper
-- around 'arrayOf'.
arrayOfRow ::
  TypeBase (ShapeBase d) NoMode ->
  d ->
  TypeBase (ShapeBase d) NoMode
arrayOfRow :: forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow TypeBase (ShapeBase d) NoMode
t d
size = TypeBase (ShapeBase d) NoMode
-> ShapeBase d -> NoMode -> TypeBase (ShapeBase d) NoMode
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf TypeBase (ShapeBase d) NoMode
t ([d] -> ShapeBase d
forall d. [d] -> ShapeBase d
Shape [d
size]) NoMode
NoMode

-- | Construct an array whose rows are the given type, and the outer
-- size is the given t'Shape'.  This is just a convenient wrapper
-- around 'arrayOf'.
arrayOfShape :: Type -> Shape -> Type
arrayOfShape :: Type -> Shape -> Type
arrayOfShape Type
t Shape
shape = Type -> Shape -> NoMode -> Type
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
arrayOf Type
t Shape
shape NoMode
NoMode

-- | Set the dimensions of an array.  If the given type is not an
-- array, return the type unchanged.
setArrayDims :: TypeBase oldshape o -> [SubExp] -> TypeBase Shape o
setArrayDims :: forall oldshape o.
TypeBase oldshape o -> [SubExp] -> TypeBase Shape o
setArrayDims TypeBase oldshape o
t [SubExp]
dims = TypeBase oldshape o
t TypeBase oldshape o -> Shape -> TypeBase Shape o
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
`setArrayShape` [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
dims

-- | Replace the size of the outermost dimension of an array.  If the
-- given type is not an array, it is returned unchanged.
setOuterSize ::
  TypeBase (ShapeBase d) m ->
  d ->
  TypeBase (ShapeBase d) m
setOuterSize :: forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
setOuterSize = Int -> TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
forall d m.
Int -> TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
setDimSize Int
0

-- | Replace the size of the given dimension of an array.  If the
-- given type is not an array, it is returned unchanged.
setDimSize ::
  Int ->
  TypeBase (ShapeBase d) m ->
  d ->
  TypeBase (ShapeBase d) m
setDimSize :: forall d m.
Int -> TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
setDimSize Int
i TypeBase (ShapeBase d) m
t d
e = TypeBase (ShapeBase d) m
t TypeBase (ShapeBase d) m -> ShapeBase d -> TypeBase (ShapeBase d) m
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
`setArrayShape` Int -> ShapeBase d -> d -> ShapeBase d
forall d. Int -> ShapeBase d -> d -> ShapeBase d
setDim Int
i (TypeBase (ShapeBase d) m -> ShapeBase d
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase (ShapeBase d) m
t) d
e

-- | Replace the outermost dimension of an array shape.
setOuterDim :: ShapeBase d -> d -> ShapeBase d
setOuterDim :: forall d. ShapeBase d -> d -> ShapeBase d
setOuterDim = Int -> ShapeBase d -> d -> ShapeBase d
forall d. Int -> ShapeBase d -> d -> ShapeBase d
setDim Int
0

-- | Replace some outermost dimensions of an array shape.
setOuterDims :: ShapeBase d -> Int -> ShapeBase d -> ShapeBase d
setOuterDims :: forall d. ShapeBase d -> Int -> ShapeBase d -> ShapeBase d
setOuterDims ShapeBase d
old Int
k ShapeBase d
new = ShapeBase d
new ShapeBase d -> ShapeBase d -> ShapeBase d
forall a. Semigroup a => a -> a -> a
<> Int -> ShapeBase d -> ShapeBase d
forall d. Int -> ShapeBase d -> ShapeBase d
stripDims Int
k ShapeBase d
old

-- | Replace the specified dimension of an array shape.
setDim :: Int -> ShapeBase d -> d -> ShapeBase d
setDim :: forall d. Int -> ShapeBase d -> d -> ShapeBase d
setDim Int
i (Shape [d]
ds) d
e = [d] -> ShapeBase d
forall d. [d] -> ShapeBase d
Shape ([d] -> ShapeBase d) -> [d] -> ShapeBase d
forall a b. (a -> b) -> a -> b
$ Int -> [d] -> [d]
forall a. Int -> [a] -> [a]
take Int
i [d]
ds [d] -> [d] -> [d]
forall a. [a] -> [a] -> [a]
++ d
e d -> [d] -> [d]
forall a. a -> [a] -> [a]
: Int -> [d] -> [d]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [d]
ds

-- | @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 Shape o -> Maybe (TypeBase Shape o)
peelArray :: forall o. Int -> TypeBase Shape o -> Maybe (TypeBase Shape o)
peelArray Int
0 TypeBase Shape o
t = TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a. a -> Maybe a
Just TypeBase Shape o
t
peelArray Int
n (Array PrimType
et Shape
shape o
o)
  | Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a. a -> Maybe a
Just (TypeBase Shape o -> Maybe (TypeBase Shape o))
-> TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
et
  | Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n = TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a. a -> Maybe a
Just (TypeBase Shape o -> Maybe (TypeBase Shape o))
-> TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> Shape -> o -> TypeBase Shape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et (Int -> Shape -> Shape
forall d. Int -> ShapeBase d -> ShapeBase d
stripDims Int
n Shape
shape) o
o
peelArray Int
_ TypeBase Shape o
_ = Maybe (TypeBase Shape o)
forall a. Maybe a
Nothing

-- | @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 (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray :: forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray Int
n (Array PrimType
et ShapeBase d
shape m
o)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ShapeBase d -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase d
shape = PrimType -> ShapeBase d -> m -> TypeBase (ShapeBase d) m
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et (Int -> ShapeBase d -> ShapeBase d
forall d. Int -> ShapeBase d -> ShapeBase d
stripDims Int
n ShapeBase d
shape) m
o
  | Bool
otherwise = PrimType -> TypeBase (ShapeBase d) m
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
et
stripArray Int
_ TypeBase (ShapeBase d) m
t = TypeBase (ShapeBase d) m
t

-- | Return the size of the given dimension.  If the dimension does
-- not exist, the zero constant is returned.
shapeSize :: Int -> Shape -> SubExp
shapeSize :: Int -> Shape -> SubExp
shapeSize Int
i Shape
shape = case Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
i ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape of
  SubExp
e : [SubExp]
_ -> SubExp
e
  [] -> Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)

-- | Return the dimensions of a type - for non-arrays, this is the
-- empty list.
arrayDims :: TypeBase (ShapeBase d) m -> [d]
arrayDims :: forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims = ShapeBase d -> [d]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase d -> [d])
-> (TypeBase (ShapeBase d) m -> ShapeBase d)
-> TypeBase (ShapeBase d) m
-> [d]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (ShapeBase d) m -> ShapeBase d
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape

-- | Return the size of the given dimension.  If the dimension does
-- not exist, the zero constant is returned.
arraySize :: Int -> TypeBase Shape o -> SubExp
arraySize :: forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
i = Int -> Shape -> SubExp
shapeSize Int
i (Shape -> SubExp)
-> (TypeBase Shape o -> Shape) -> TypeBase Shape o -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape o -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape

-- | Return the size of the given dimension in the first element of
-- the given type list.  If the dimension does not exist, or no types
-- are given, the zero constant is returned.
arraysSize :: Int -> [TypeBase Shape o] -> SubExp
arraysSize :: forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
_ [] = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
arraysSize Int
i (TypeBase Shape o
t : [TypeBase Shape o]
_) = Int -> TypeBase Shape o -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
i TypeBase Shape o
t

-- | Return the immediate row-type of an array.  For @[][]t@, this
-- would be @[]t@.
rowType :: TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType :: forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType = Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray Int
1

-- | A type is a primitive type if it is not an array or memory block.
primType :: TypeBase shape o -> Bool
primType :: forall shape o. TypeBase shape o -> Bool
primType Prim {} = Bool
True
primType TypeBase shape o
_ = Bool
False

-- | Is this an accumulator?
isAcc :: TypeBase shape o -> Bool
isAcc :: forall shape o. TypeBase shape o -> Bool
isAcc Acc {} = Bool
True
isAcc TypeBase shape o
_ = Bool
False

-- | Is this a memory block?
isMem :: TypeBase shape o -> Bool
isMem :: forall shape o. TypeBase shape o -> Bool
isMem Mem {} = Bool
True
isMem TypeBase shape o
_ = Bool
False

-- | Returns the bottommost type of an array.  For @[][]i32@, this
-- would be @i32@.  If the given type is not an array, it is returned.
elemType :: TypeBase shape o -> PrimType
elemType :: forall shape o. TypeBase shape o -> PrimType
elemType (Array PrimType
t shape
_ o
_) = PrimType
t
elemType (Prim PrimType
t) = PrimType
t
elemType Acc {} = [Char] -> PrimType
forall a. HasCallStack => [Char] -> a
error [Char]
"elemType Acc"
elemType Mem {} = [Char] -> PrimType
forall a. HasCallStack => [Char] -> a
error [Char]
"elemType Mem"

-- | Swap the two outer dimensions of the type.
transposeType :: Type -> Type
transposeType :: Type -> Type
transposeType = [Int] -> Type -> Type
forall d m.
[Int] -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rearrangeType [Int
1, Int
0]

-- | Rearrange the dimensions of the type.  If the length of the
-- permutation does not match the rank of the type, the permutation
-- will be extended with identity.
rearrangeType :: [Int] -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rearrangeType :: forall d m.
[Int] -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rearrangeType [Int]
perm TypeBase (ShapeBase d) m
t =
  TypeBase (ShapeBase d) m
t TypeBase (ShapeBase d) m -> ShapeBase d -> TypeBase (ShapeBase d) m
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
`setArrayShape` [d] -> ShapeBase d
forall d. [d] -> ShapeBase d
Shape ([Int] -> [d] -> [d]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm' ([d] -> [d]) -> [d] -> [d]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase d) m -> [d]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase (ShapeBase d) m
t)
  where
    perm' :: [Int]
perm' = [Int]
perm [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
perm .. TypeBase (ShapeBase d) m -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank TypeBase (ShapeBase d) m
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

-- | Transform any t'SubExp's in the type.
mapOnExtType ::
  (Monad m) =>
  (SubExp -> m SubExp) ->
  TypeBase ExtShape o ->
  m (TypeBase ExtShape o)
mapOnExtType :: forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp)
-> TypeBase ExtShape o -> m (TypeBase ExtShape o)
mapOnExtType SubExp -> m SubExp
_ (Prim PrimType
bt) =
  TypeBase ExtShape o -> m (TypeBase ExtShape o)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase ExtShape o -> m (TypeBase ExtShape o))
-> TypeBase ExtShape o -> m (TypeBase ExtShape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase ExtShape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
bt
mapOnExtType SubExp -> m SubExp
f (Acc VName
acc Shape
ispace [Type]
ts) =
  VName -> Shape -> [Type] -> TypeBase ExtShape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc (VName -> Shape -> [Type] -> TypeBase ExtShape o)
-> m VName -> m (Shape -> [Type] -> TypeBase ExtShape o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m VName
f' VName
acc m (Shape -> [Type] -> TypeBase ExtShape o)
-> m Shape -> m ([Type] -> TypeBase ExtShape o)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> Shape -> m Shape
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) -> ShapeBase a -> f (ShapeBase b)
traverse SubExp -> m SubExp
f Shape
ispace m ([Type] -> TypeBase ExtShape o)
-> m [Type] -> m (TypeBase ExtShape o)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> [Type] -> m [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> m SubExp) -> Type -> m Type
forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape o -> m (TypeBase Shape o)
mapOnType SubExp -> m SubExp
f) [Type]
ts
  where
    f' :: VName -> m VName
f' VName
v = do
      x <- SubExp -> m SubExp
f (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
      case x of
        Var VName
v' -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v'
        Constant {} -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
mapOnExtType SubExp -> m SubExp
_ (Mem Space
space) =
  TypeBase ExtShape o -> m (TypeBase ExtShape o)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase ExtShape o -> m (TypeBase ExtShape o))
-> TypeBase ExtShape o -> m (TypeBase ExtShape o)
forall a b. (a -> b) -> a -> b
$ Space -> TypeBase ExtShape o
forall shape o. Space -> TypeBase shape o
Mem Space
space
mapOnExtType SubExp -> m SubExp
f (Array PrimType
t ExtShape
shape o
o) =
  PrimType -> ExtShape -> o -> TypeBase ExtShape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
t
    (ExtShape -> o -> TypeBase ExtShape o)
-> m ExtShape -> m (o -> TypeBase ExtShape o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Ext SubExp] -> ExtShape
forall d. [d] -> ShapeBase d
Shape ([Ext SubExp] -> ExtShape) -> m [Ext SubExp] -> m ExtShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ext SubExp -> m (Ext SubExp)) -> [Ext SubExp] -> m [Ext SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> m SubExp) -> Ext SubExp -> m (Ext SubExp)
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) -> Ext a -> f (Ext b)
traverse SubExp -> m SubExp
f) (ExtShape -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims ExtShape
shape))
    m (o -> TypeBase ExtShape o) -> m o -> m (TypeBase ExtShape o)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> o -> m o
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure o
o

-- | Transform any t'SubExp's in the type.
mapOnType ::
  (Monad m) =>
  (SubExp -> m SubExp) ->
  TypeBase Shape o ->
  m (TypeBase Shape o)
mapOnType :: forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape o -> m (TypeBase Shape o)
mapOnType SubExp -> m SubExp
_ (Prim PrimType
bt) = TypeBase Shape o -> m (TypeBase Shape o)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Shape o -> m (TypeBase Shape o))
-> TypeBase Shape o -> m (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
bt
mapOnType SubExp -> m SubExp
f (Acc VName
acc Shape
ispace [Type]
ts) =
  VName -> Shape -> [Type] -> TypeBase Shape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc (VName -> Shape -> [Type] -> TypeBase Shape o)
-> m VName -> m (Shape -> [Type] -> TypeBase Shape o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m VName
f' VName
acc m (Shape -> [Type] -> TypeBase Shape o)
-> m Shape -> m ([Type] -> TypeBase Shape o)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (SubExp -> m SubExp) -> Shape -> m Shape
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) -> ShapeBase a -> f (ShapeBase b)
traverse SubExp -> m SubExp
f Shape
ispace m ([Type] -> TypeBase Shape o) -> m [Type] -> m (TypeBase Shape o)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> [Type] -> m [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SubExp -> m SubExp) -> Type -> m Type
forall (m :: * -> *) o.
Monad m =>
(SubExp -> m SubExp) -> TypeBase Shape o -> m (TypeBase Shape o)
mapOnType SubExp -> m SubExp
f) [Type]
ts
  where
    f' :: VName -> m VName
f' VName
v = do
      x <- SubExp -> m SubExp
f (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
      case x of
        Var VName
v' -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v'
        Constant {} -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
mapOnType SubExp -> m SubExp
_ (Mem Space
space) = TypeBase Shape o -> m (TypeBase Shape o)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Shape o -> m (TypeBase Shape o))
-> TypeBase Shape o -> m (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ Space -> TypeBase Shape o
forall shape o. Space -> TypeBase shape o
Mem Space
space
mapOnType SubExp -> m SubExp
f (Array PrimType
t Shape
shape o
o) =
  PrimType -> Shape -> o -> TypeBase Shape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
t
    (Shape -> o -> TypeBase Shape o)
-> m Shape -> m (o -> TypeBase Shape o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> m [SubExp] -> m Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> m SubExp
f (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape))
    m (o -> TypeBase Shape o) -> m o -> m (TypeBase Shape o)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> o -> m o
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure o
o

-- | @diet t@ returns a description of how a function parameter of
-- type @t@ might consume its argument.  An accumulator is consumed by
-- any use, so it is unconditionally 'Consume'.
diet :: TypeBase shape Diet -> Diet
diet :: forall shape. TypeBase shape Diet -> Diet
diet Prim {} = Diet
Observe
diet Acc {} = Diet
Consume
diet (Array PrimType
_ shape
_ Diet
o) = Diet
o
diet Mem {} = Diet
Observe

-- | Add the given 'Diet' to the types.
toDecl ::
  TypeBase shape NoMode ->
  Diet ->
  TypeBase shape Diet
toDecl :: forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
toDecl (Prim PrimType
t) Diet
_ = PrimType -> TypeBase shape Diet
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t
toDecl (Acc VName
acc Shape
ispace [Type]
ts) Diet
_ = VName -> Shape -> [Type] -> TypeBase shape Diet
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
toDecl (Array PrimType
et shape
shape NoMode
_) Diet
o = PrimType -> shape -> Diet -> TypeBase shape Diet
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et shape
shape Diet
o
toDecl (Mem Space
space) Diet
_ = Space -> TypeBase shape Diet
forall shape o. Space -> TypeBase shape o
Mem Space
space

-- | Remove the mode from the type.
fromDecl ::
  TypeBase shape o ->
  TypeBase shape NoMode
fromDecl :: forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl (Prim PrimType
t) = PrimType -> TypeBase shape NoMode
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
t
fromDecl (Acc VName
acc Shape
ispace [Type]
ts) = VName -> Shape -> [Type] -> TypeBase shape NoMode
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
fromDecl (Array PrimType
et shape
shape o
_) = PrimType -> shape -> NoMode -> TypeBase shape NoMode
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
et shape
shape NoMode
NoMode
fromDecl (Mem Space
space) = Space -> TypeBase shape NoMode
forall shape o. Space -> TypeBase shape o
Mem Space
space

-- | If an existential, then return its existential index.
isExt :: Ext a -> Maybe Int
isExt :: forall a. Ext a -> Maybe Int
isExt (Ext Int
i) = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i
isExt Ext a
_ = Maybe Int
forall a. Maybe a
Nothing

-- | If a known size, then return that size.
isFree :: Ext a -> Maybe a
isFree :: forall a. Ext a -> Maybe a
isFree (Free a
d) = a -> Maybe a
forall a. a -> Maybe a
Just a
d
isFree Ext a
_ = Maybe a
forall a. Maybe a
Nothing

-- | Given the existential return type of a function, and the shapes
-- of the values returned by the function, return the existential
-- shape context.  That is, those sizes that are existential in the
-- return type.
extractShapeContext :: [TypeBase ExtShape o] -> [[a]] -> [a]
extractShapeContext :: forall o a. [TypeBase ExtShape o] -> [[a]] -> [a]
extractShapeContext [TypeBase ExtShape o]
ts [[a]]
shapes =
  State (Set Int) [a] -> Set Int -> [a]
forall s a. State s a -> s -> a
evalState ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a])
-> StateT (Set Int) Identity [[a]] -> State (Set Int) [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase ExtShape o -> [a] -> State (Set Int) [a])
-> [TypeBase ExtShape o]
-> [[a]]
-> StateT (Set Int) Identity [[a]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase ExtShape o -> [a] -> State (Set Int) [a]
forall {f :: * -> *} {a} {o} {a}.
MonadState (Set Int) f =>
TypeBase (ShapeBase (Ext a)) o -> [a] -> f [a]
extract [TypeBase ExtShape o]
ts [[a]]
shapes) Set Int
forall a. Set a
S.empty
  where
    extract :: TypeBase (ShapeBase (Ext a)) o -> [a] -> f [a]
extract TypeBase (ShapeBase (Ext a)) o
t [a]
shape =
      [Maybe a] -> [a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe a] -> [a]) -> f [Maybe a] -> f [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ext a -> a -> f (Maybe a)) -> [Ext a] -> [a] -> f [Maybe a]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Ext a -> a -> f (Maybe a)
forall {m :: * -> *} {a} {a}.
MonadState (Set Int) m =>
Ext a -> a -> m (Maybe a)
extract' (ShapeBase (Ext a) -> [Ext a]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase (Ext a) -> [Ext a]) -> ShapeBase (Ext a) -> [Ext a]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext a)) o -> ShapeBase (Ext a)
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase (ShapeBase (Ext a)) o
t) [a]
shape
    extract' :: Ext a -> a -> m (Maybe a)
extract' (Ext Int
x) a
v = do
      seen <- (Set Int -> Bool) -> m Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Set Int -> Bool) -> m Bool) -> (Set Int -> Bool) -> m Bool
forall a b. (a -> b) -> a -> b
$ Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Int
x
      if seen
        then pure Nothing
        else do
          modify $ S.insert x
          pure $ Just v
    extract' (Free a
_) a
_ = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

-- | The 'Ext' integers used for existential sizes in the given types.
shapeContext :: [TypeBase ExtShape o] -> S.Set Int
shapeContext :: forall o. [TypeBase ExtShape o] -> Set Int
shapeContext = [Int] -> Set Int
forall a. Ord a => [a] -> Set a
S.fromList ([Int] -> Set Int)
-> ([TypeBase ExtShape o] -> [Int])
-> [TypeBase ExtShape o]
-> Set Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase ExtShape o -> [Int]) -> [TypeBase ExtShape o] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Ext SubExp -> Maybe Int) -> [Ext SubExp] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Ext SubExp -> Maybe Int
forall a. Ext a -> Maybe Int
isExt ([Ext SubExp] -> [Int])
-> (TypeBase ExtShape o -> [Ext SubExp])
-> TypeBase ExtShape o
-> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtShape -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ExtShape -> [Ext SubExp])
-> (TypeBase ExtShape o -> ExtShape)
-> TypeBase ExtShape o
-> [Ext SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape o -> ExtShape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape)

-- | If all dimensions of the given 'ExtShape' are statically known,
-- change to the corresponding t'Shape'.
hasStaticShape :: TypeBase ExtShape o -> Maybe (TypeBase Shape o)
hasStaticShape :: forall o. TypeBase ExtShape o -> Maybe (TypeBase Shape o)
hasStaticShape (Prim PrimType
bt) = TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a. a -> Maybe a
Just (TypeBase Shape o -> Maybe (TypeBase Shape o))
-> TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
bt
hasStaticShape (Acc VName
acc Shape
ispace [Type]
ts) = TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a. a -> Maybe a
Just (TypeBase Shape o -> Maybe (TypeBase Shape o))
-> TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ VName -> Shape -> [Type] -> TypeBase Shape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts
hasStaticShape (Mem Space
space) = TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a. a -> Maybe a
Just (TypeBase Shape o -> Maybe (TypeBase Shape o))
-> TypeBase Shape o -> Maybe (TypeBase Shape o)
forall a b. (a -> b) -> a -> b
$ Space -> TypeBase Shape o
forall shape o. Space -> TypeBase shape o
Mem Space
space
hasStaticShape (Array PrimType
bt (Shape [Ext SubExp]
shape) o
o) =
  PrimType -> Shape -> o -> TypeBase Shape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
bt (Shape -> o -> TypeBase Shape o)
-> Maybe Shape -> Maybe (o -> TypeBase Shape o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> Maybe [SubExp] -> Maybe Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ext SubExp -> Maybe SubExp) -> [Ext SubExp] -> Maybe [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Ext SubExp -> Maybe SubExp
forall a. Ext a -> Maybe a
isFree [Ext SubExp]
shape) Maybe (o -> TypeBase Shape o)
-> Maybe o -> Maybe (TypeBase Shape o)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> o -> Maybe o
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure o
o

-- | Given two lists of 'ExtType's of the same length, return a list of
-- 'ExtType's that generalises the two operands.
generaliseExtTypes ::
  [TypeBase ExtShape o] ->
  [TypeBase ExtShape o] ->
  [TypeBase ExtShape o]
generaliseExtTypes :: forall o.
[TypeBase ExtShape o]
-> [TypeBase ExtShape o] -> [TypeBase ExtShape o]
generaliseExtTypes [TypeBase ExtShape o]
rt1 [TypeBase ExtShape o]
rt2 =
  State (Int, Map Int Int) [TypeBase ExtShape o]
-> (Int, Map Int Int) -> [TypeBase ExtShape o]
forall s a. State s a -> s -> a
evalState ((TypeBase ExtShape o
 -> TypeBase ExtShape o
 -> StateT (Int, Map Int Int) Identity (TypeBase ExtShape o))
-> [TypeBase ExtShape o]
-> [TypeBase ExtShape o]
-> State (Int, Map Int Int) [TypeBase ExtShape o]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase ExtShape o
-> TypeBase ExtShape o
-> StateT (Int, Map Int Int) Identity (TypeBase ExtShape o)
forall {f :: * -> *} {a} {o} {o}.
(MonadState (Int, Map Int Int) f, Eq a) =>
TypeBase (ShapeBase (Ext a)) o
-> TypeBase (ShapeBase (Ext a)) o
-> f (TypeBase (ShapeBase (Ext a)) o)
unifyExtShapes [TypeBase ExtShape o]
rt1 [TypeBase ExtShape o]
rt2) (Int
0, Map Int Int
forall k a. Map k a
M.empty)
  where
    unifyExtShapes :: TypeBase (ShapeBase (Ext a)) o
-> TypeBase (ShapeBase (Ext a)) o
-> f (TypeBase (ShapeBase (Ext a)) o)
unifyExtShapes TypeBase (ShapeBase (Ext a)) o
t1 TypeBase (ShapeBase (Ext a)) o
t2 =
      TypeBase (ShapeBase (Ext a)) o
-> ShapeBase (Ext a) -> TypeBase (ShapeBase (Ext a)) o
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
setArrayShape TypeBase (ShapeBase (Ext a)) o
t1 (ShapeBase (Ext a) -> TypeBase (ShapeBase (Ext a)) o)
-> ([Ext a] -> ShapeBase (Ext a))
-> [Ext a]
-> TypeBase (ShapeBase (Ext a)) o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Ext a] -> ShapeBase (Ext a)
forall d. [d] -> ShapeBase d
Shape
        ([Ext a] -> TypeBase (ShapeBase (Ext a)) o)
-> f [Ext a] -> f (TypeBase (ShapeBase (Ext a)) o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ext a -> Ext a -> f (Ext a)) -> [Ext a] -> [Ext a] -> f [Ext a]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM
          Ext a -> Ext a -> f (Ext a)
forall {f :: * -> *} {a}.
(MonadState (Int, Map Int Int) f, Eq a) =>
Ext a -> Ext a -> f (Ext a)
unifyExtDims
          (ShapeBase (Ext a) -> [Ext a]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase (Ext a) -> [Ext a]) -> ShapeBase (Ext a) -> [Ext a]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext a)) o -> ShapeBase (Ext a)
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase (ShapeBase (Ext a)) o
t1)
          (ShapeBase (Ext a) -> [Ext a]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase (Ext a) -> [Ext a]) -> ShapeBase (Ext a) -> [Ext a]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext a)) o -> ShapeBase (Ext a)
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase (ShapeBase (Ext a)) o
t2)
    unifyExtDims :: Ext a -> Ext a -> f (Ext a)
unifyExtDims (Free a
se1) (Free a
se2)
      | a
se1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
se2 = Ext a -> f (Ext a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ext a -> f (Ext a)) -> Ext a -> f (Ext a)
forall a b. (a -> b) -> a -> b
$ a -> Ext a
forall a. a -> Ext a
Free a
se1 -- Arbitrary
      | Bool
otherwise = do
          (n, m) <- f (Int, Map Int Int)
forall s (m :: * -> *). MonadState s m => m s
get
          put (n + 1, m)
          pure $ Ext n
    unifyExtDims (Ext Int
x) (Ext Int
y)
      | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y = Int -> Ext a
forall a. Int -> Ext a
Ext (Int -> Ext a) -> f Int -> f (Ext a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f Int -> (Int -> f Int) -> Maybe Int -> f Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> f Int
forall {m :: * -> *} {b} {k}.
(MonadState (b, Map k b) m, Num b, Ord k) =>
k -> m b
new Int
x) Int -> f Int
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> f Int) -> f (Maybe Int) -> f Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Int, Map Int Int) -> Maybe Int) -> f (Maybe Int)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Int -> Map Int Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
x (Map Int Int -> Maybe Int)
-> ((Int, Map Int Int) -> Map Int Int)
-> (Int, Map Int Int)
-> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Map Int Int) -> Map Int Int
forall a b. (a, b) -> b
snd))
    unifyExtDims (Ext Int
x) Ext a
_ = Int -> Ext a
forall a. Int -> Ext a
Ext (Int -> Ext a) -> f Int -> f (Ext a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f Int
forall {m :: * -> *} {b} {k}.
(MonadState (b, Map k b) m, Num b, Ord k) =>
k -> m b
new Int
x
    unifyExtDims Ext a
_ (Ext Int
x) = Int -> Ext a
forall a. Int -> Ext a
Ext (Int -> Ext a) -> f Int -> f (Ext a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f Int
forall {m :: * -> *} {b} {k}.
(MonadState (b, Map k b) m, Num b, Ord k) =>
k -> m b
new Int
x
    new :: k -> m b
new k
x = do
      (n, m) <- m (b, Map k b)
forall s (m :: * -> *). MonadState s m => m s
get
      put (n + 1, M.insert x n m)
      pure n

-- | Given a list of 'ExtType's and a list of "forbidden" names,
-- modify the dimensions of the 'ExtType's such that they are 'Ext'
-- where they were previously 'Free' with a variable in the set of
-- forbidden names.
existentialiseExtTypes :: [VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes :: [VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes [VName]
inaccessible = (ExtType -> ExtType) -> [ExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map ExtType -> ExtType
forall {o}. TypeBase ExtShape o -> TypeBase ExtShape o
makeBoundShapesFree
  where
    makeBoundShapesFree :: TypeBase ExtShape o -> TypeBase ExtShape o
makeBoundShapesFree =
      (ExtShape -> ExtShape)
-> TypeBase ExtShape o -> TypeBase ExtShape o
forall newshape oldshape o.
ArrayShape newshape =>
(oldshape -> newshape)
-> TypeBase oldshape o -> TypeBase newshape o
modifyArrayShape ((ExtShape -> ExtShape)
 -> TypeBase ExtShape o -> TypeBase ExtShape o)
-> (ExtShape -> ExtShape)
-> TypeBase ExtShape o
-> TypeBase ExtShape o
forall a b. (a -> b) -> a -> b
$ (Ext SubExp -> Ext SubExp) -> ExtShape -> ExtShape
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ext SubExp -> Ext SubExp
checkDim
    checkDim :: Ext SubExp -> Ext SubExp
checkDim (Free (Var VName
v))
      | Just Int
i <- VName
v VName -> [VName] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`L.elemIndex` [VName]
inaccessible =
          Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i
    checkDim Ext SubExp
d = Ext SubExp
d

-- | Produce a mapping for the dimensions context.
shapeExtMapping :: [TypeBase ExtShape o] -> [TypeBase Shape o1] -> M.Map Int SubExp
shapeExtMapping :: forall o o1.
[TypeBase ExtShape o] -> [TypeBase Shape o1] -> Map Int SubExp
shapeExtMapping = (TypeBase ExtShape o -> [Ext SubExp])
-> (TypeBase Shape o1 -> [SubExp])
-> (Ext SubExp -> SubExp -> Map Int SubExp)
-> (Map Int SubExp -> Map Int SubExp -> Map Int SubExp)
-> [TypeBase ExtShape o]
-> [TypeBase Shape o1]
-> Map Int SubExp
forall res t1 dim1 t2 dim2.
Monoid res =>
(t1 -> [dim1])
-> (t2 -> [dim2])
-> (dim1 -> dim2 -> res)
-> (res -> res -> res)
-> [t1]
-> [t2]
-> res
dimMapping TypeBase ExtShape o -> [Ext SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase Shape o1 -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Ext SubExp -> SubExp -> Map Int SubExp
forall {a} {a}. Ext a -> a -> Map Int a
match Map Int SubExp -> Map Int SubExp -> Map Int SubExp
forall a. Monoid a => a -> a -> a
mappend
  where
    match :: Ext a -> a -> Map Int a
match Free {} a
_ = Map Int a
forall a. Monoid a => a
mempty
    match (Ext Int
i) a
dim = Int -> a -> Map Int a
forall k a. k -> a -> Map k a
M.singleton Int
i a
dim

dimMapping ::
  (Monoid res) =>
  (t1 -> [dim1]) ->
  (t2 -> [dim2]) ->
  (dim1 -> dim2 -> res) ->
  (res -> res -> res) ->
  [t1] ->
  [t2] ->
  res
dimMapping :: forall res t1 dim1 t2 dim2.
Monoid res =>
(t1 -> [dim1])
-> (t2 -> [dim2])
-> (dim1 -> dim2 -> res)
-> (res -> res -> res)
-> [t1]
-> [t2]
-> res
dimMapping t1 -> [dim1]
getDims1 t2 -> [dim2]
getDims2 dim1 -> dim2 -> res
f res -> res -> res
comb [t1]
ts1 [t2]
ts2 =
  (res -> res -> res) -> res -> [res] -> res
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' res -> res -> res
comb res
forall a. Monoid a => a
mempty ([res] -> res) -> [res] -> res
forall a b. (a -> b) -> a -> b
$ [[res]] -> [res]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[res]] -> [res]) -> [[res]] -> [res]
forall a b. (a -> b) -> a -> b
$ ([dim1] -> [dim2] -> [res]) -> [[dim1]] -> [[dim2]] -> [[res]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((dim1 -> dim2 -> res) -> [dim1] -> [dim2] -> [res]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith dim1 -> dim2 -> res
f) ((t1 -> [dim1]) -> [t1] -> [[dim1]]
forall a b. (a -> b) -> [a] -> [b]
map t1 -> [dim1]
getDims1 [t1]
ts1) ((t2 -> [dim2]) -> [t2] -> [[dim2]]
forall a b. (a -> b) -> [a] -> [b]
map t2 -> [dim2]
getDims2 [t2]
ts2)

-- | @IntType Int8@
int8 :: PrimType
int8 :: PrimType
int8 = IntType -> PrimType
IntType IntType
Int8

-- | @IntType Int16@
int16 :: PrimType
int16 :: PrimType
int16 = IntType -> PrimType
IntType IntType
Int16

-- | @IntType Int32@
int32 :: PrimType
int32 :: PrimType
int32 = IntType -> PrimType
IntType IntType
Int32

-- | @IntType Int64@
int64 :: PrimType
int64 :: PrimType
int64 = IntType -> PrimType
IntType IntType
Int64

-- | @FloatType Float32@
float32 :: PrimType
float32 :: PrimType
float32 = FloatType -> PrimType
FloatType FloatType
Float32

-- | @FloatType Float64@
float64 :: PrimType
float64 :: PrimType
float64 = FloatType -> PrimType
FloatType FloatType
Float64

-- | Typeclass for things that contain 'Type's.
class Typed t where
  typeOf :: t -> Type

instance Typed Type where
  typeOf :: Type -> Type
typeOf = Type -> Type
forall a. a -> a
id

instance Typed DeclType where
  typeOf :: DeclType -> Type
typeOf = DeclType -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl

instance Typed Ident where
  typeOf :: Ident -> Type
typeOf = Ident -> Type
identType

instance (Typed dec) => Typed (Param dec) where
  typeOf :: Param dec -> Type
typeOf = dec -> Type
forall t. Typed t => t -> Type
typeOf (dec -> Type) -> (Param dec -> dec) -> Param dec -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param dec -> dec
forall dec. Param dec -> dec
paramDec

instance (Typed dec) => Typed (PatElem dec) where
  typeOf :: PatElem dec -> Type
typeOf = dec -> Type
forall t. Typed t => t -> Type
typeOf (dec -> Type) -> (PatElem dec -> dec) -> PatElem dec -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem dec -> dec
forall dec. PatElem dec -> dec
patElemDec

instance (Typed b) => Typed (a, b) where
  typeOf :: (a, b) -> Type
typeOf = b -> Type
forall t. Typed t => t -> Type
typeOf (b -> Type) -> ((a, b) -> b) -> (a, b) -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> b
forall a b. (a, b) -> b
snd

-- | Typeclass for things that contain 'DeclType's.
class DeclTyped t where
  declTypeOf :: t -> DeclType

instance DeclTyped DeclType where
  declTypeOf :: DeclType -> DeclType
declTypeOf = DeclType -> DeclType
forall a. a -> a
id

instance (DeclTyped dec) => DeclTyped (Param dec) where
  declTypeOf :: Param dec -> DeclType
declTypeOf = dec -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf (dec -> DeclType) -> (Param dec -> dec) -> Param dec -> DeclType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param dec -> dec
forall dec. Param dec -> dec
paramDec

-- | Typeclass for things that contain 'ExtType's.
class (FixExt t) => ExtTyped t where
  extTypeOf :: t -> ExtType

instance ExtTyped ExtType where
  extTypeOf :: ExtType -> ExtType
extTypeOf = ExtType -> ExtType
forall a. a -> a
id

instance ExtTyped DeclExtType where
  extTypeOf :: DeclExtType -> ExtType
extTypeOf = DeclExtType -> ExtType
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl

-- | Something with an existential context that can be (partially)
-- fixed.
class FixExt t where
  -- | Fix the given existentional variable to the indicated free
  -- value.
  fixExt :: Int -> SubExp -> t -> t

  -- | Map a function onto any existential.
  mapExt :: (Int -> Int) -> t -> t

instance (FixExt shape, ArrayShape shape) => FixExt (TypeBase shape o) where
  fixExt :: Int -> SubExp -> TypeBase shape o -> TypeBase shape o
fixExt Int
i SubExp
se = (shape -> shape) -> TypeBase shape o -> TypeBase shape o
forall newshape oldshape o.
ArrayShape newshape =>
(oldshape -> newshape)
-> TypeBase oldshape o -> TypeBase newshape o
modifyArrayShape ((shape -> shape) -> TypeBase shape o -> TypeBase shape o)
-> (shape -> shape) -> TypeBase shape o -> TypeBase shape o
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> shape -> shape
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se
  mapExt :: (Int -> Int) -> TypeBase shape o -> TypeBase shape o
mapExt Int -> Int
f = (shape -> shape) -> TypeBase shape o -> TypeBase shape o
forall newshape oldshape o.
ArrayShape newshape =>
(oldshape -> newshape)
-> TypeBase oldshape o -> TypeBase newshape o
modifyArrayShape ((shape -> shape) -> TypeBase shape o -> TypeBase shape o)
-> (shape -> shape) -> TypeBase shape o -> TypeBase shape o
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> shape -> shape
forall t. FixExt t => (Int -> Int) -> t -> t
mapExt Int -> Int
f

instance (FixExt d) => FixExt (ShapeBase d) where
  fixExt :: Int -> SubExp -> ShapeBase d -> ShapeBase d
fixExt Int
i SubExp
se = (d -> d) -> ShapeBase d -> ShapeBase d
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((d -> d) -> ShapeBase d -> ShapeBase d)
-> (d -> d) -> ShapeBase d -> ShapeBase d
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> d -> d
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se
  mapExt :: (Int -> Int) -> ShapeBase d -> ShapeBase d
mapExt Int -> Int
f = (d -> d) -> ShapeBase d -> ShapeBase d
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((d -> d) -> ShapeBase d -> ShapeBase d)
-> (d -> d) -> ShapeBase d -> ShapeBase d
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> d -> d
forall t. FixExt t => (Int -> Int) -> t -> t
mapExt Int -> Int
f

instance (FixExt a) => FixExt [a] where
  fixExt :: Int -> SubExp -> [a] -> [a]
fixExt Int
i SubExp
se = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a) -> [a] -> [a]) -> (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> a -> a
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se
  mapExt :: (Int -> Int) -> [a] -> [a]
mapExt Int -> Int
f = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a) -> [a] -> [a]) -> (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> a -> a
forall t. FixExt t => (Int -> Int) -> t -> t
mapExt Int -> Int
f

instance FixExt ExtSize where
  fixExt :: Int -> SubExp -> Ext SubExp -> Ext SubExp
fixExt Int
i SubExp
se (Ext Int
j)
    | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
i = Int -> Ext SubExp
forall a. Int -> Ext a
Ext (Int -> Ext SubExp) -> Int -> Ext SubExp
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i = SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
se
    | Bool
otherwise = Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
j
  fixExt Int
_ SubExp
_ (Free SubExp
x) = SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
x

  mapExt :: (Int -> Int) -> Ext SubExp -> Ext SubExp
mapExt Int -> Int
f (Ext Int
i) = Int -> Ext SubExp
forall a. Int -> Ext a
Ext (Int -> Ext SubExp) -> Int -> Ext SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Int
f Int
i
  mapExt Int -> Int
_ (Free SubExp
x) = SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
x

instance FixExt () where
  fixExt :: Int -> SubExp -> () -> ()
fixExt Int
_ SubExp
_ () = ()
  mapExt :: (Int -> Int) -> () -> ()
mapExt Int -> Int
_ () = ()