{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.IR.Prop.TypeOf
( expExtType,
expExtTypeSize,
subExpType,
primOpType,
mapType,
module Futhark.IR.RetType,
module Futhark.IR.Prop.Scope,
TypedOp (..),
)
where
import Data.Maybe
import Futhark.IR.Prop.Constants
import Futhark.IR.Prop.Patterns
import Futhark.IR.Prop.Reshape
import Futhark.IR.Prop.Scope
import Futhark.IR.Prop.Types
import Futhark.IR.RetType
import Futhark.IR.Syntax
subExpType :: HasScope t m => SubExp -> m Type
subExpType :: forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType (Constant PrimValue
val) = Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
subExpType (Var VName
name) = VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
name
mapType :: SubExp -> Lambda lore -> [Type]
mapType :: forall lore. SubExp -> Lambda lore -> [Type]
mapType SubExp
outersize Lambda lore
f =
[ Type -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf Type
t ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
outersize]) NoUniqueness
NoUniqueness
| Type
t <- Lambda lore -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda lore
f
]
primOpType :: HasScope lore m => BasicOp -> m [Type]
primOpType :: forall lore (m :: * -> *). HasScope lore m => BasicOp -> m [Type]
primOpType (SubExp SubExp
se) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
primOpType (Opaque SubExp
se) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
primOpType (ArrayLit [SubExp]
es Type
rt) =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf Type
rt ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
n]) NoUniqueness
NoUniqueness]
where
n :: SubExp
n = IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
es
primOpType (BinOp BinOp
bop SubExp
_ SubExp
_) =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ BinOp -> PrimType
binOpType BinOp
bop]
primOpType (UnOp UnOp
_ SubExp
x) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
x
primOpType CmpOp {} =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Bool]
primOpType (ConvOp ConvOp
conv SubExp
_) =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ (PrimType, PrimType) -> PrimType
forall a b. (a, b) -> b
snd ((PrimType, PrimType) -> PrimType)
-> (PrimType, PrimType) -> PrimType
forall a b. (a -> b) -> a -> b
$ ConvOp -> (PrimType, PrimType)
convOpType ConvOp
conv]
primOpType (Index VName
ident Slice SubExp
slice) =
Type -> [Type]
result (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
ident
where
result :: Type -> [Type]
result Type
t = [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t) Type -> Shape -> Type
`arrayOfShape` Shape
shape]
shape :: Shape
shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ (DimIndex SubExp -> Maybe SubExp) -> Slice SubExp -> [SubExp]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe DimIndex SubExp -> Maybe SubExp
forall {a}. DimIndex a -> Maybe a
dimSize Slice SubExp
slice
dimSize :: DimIndex a -> Maybe a
dimSize (DimSlice a
_ a
d a
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
d
dimSize DimFix {} = Maybe a
forall a. Maybe a
Nothing
primOpType (Update VName
src Slice SubExp
_ SubExp
_) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
src
primOpType (Iota SubExp
n SubExp
_ SubExp
_ IntType
et) =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase Shape Any -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf (PrimType -> TypeBase Shape Any
forall shape u. PrimType -> TypeBase shape u
Prim (IntType -> PrimType
IntType IntType
et)) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
n]) NoUniqueness
NoUniqueness]
primOpType (Replicate (Shape []) SubExp
e) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e
primOpType (Replicate Shape
shape SubExp
e) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> (Type -> Type) -> Type -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type -> Shape -> Type) -> Shape -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Shape -> Type
arrayOfShape Shape
shape (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e
primOpType (Scratch PrimType
t [SubExp]
shape) =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase Shape Any -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf (PrimType -> TypeBase Shape Any
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
shape) NoUniqueness
NoUniqueness]
primOpType (Reshape [] VName
e) =
Type -> [Type]
forall {shape} {u} {shape} {u}.
TypeBase shape u -> [TypeBase shape u]
result (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
e
where
result :: TypeBase shape u -> [TypeBase shape u]
result TypeBase shape u
t = [PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> TypeBase shape u) -> PrimType -> TypeBase shape u
forall a b. (a -> b) -> a -> b
$ TypeBase shape u -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType TypeBase shape u
t]
primOpType (Reshape [DimChange SubExp]
shape VName
e) =
Type -> [Type]
result (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
e
where
result :: Type -> [Type]
result Type
t = [Type
t Type -> Shape -> Type
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`setArrayShape` [DimChange SubExp] -> Shape
newShape [DimChange SubExp]
shape]
primOpType (Rearrange [Int]
perm VName
e) =
Type -> [Type]
result (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
e
where
result :: Type -> [Type]
result Type
t = [[Int] -> Type -> Type
rearrangeType [Int]
perm Type
t]
primOpType (Rotate [SubExp]
_ VName
e) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
e
primOpType (Concat Int
i VName
x [VName]
_ SubExp
ressize) =
Type -> [Type]
result (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
x
where
result :: Type -> [Type]
result Type
xt = [Int -> Type -> SubExp -> Type
forall d u.
ArrayShape (ShapeBase d) =>
Int -> TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u
setDimSize Int
i Type
xt SubExp
ressize]
primOpType (Copy VName
v) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
primOpType (Manifest [Int]
_ VName
v) =
Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> [Type]) -> m Type -> m [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
primOpType Assert {} =
[Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Cert]
expExtType ::
(HasScope lore m, TypedOp (Op lore)) =>
Exp lore ->
m [ExtType]
expExtType :: forall lore (m :: * -> *).
(HasScope lore m, TypedOp (Op lore)) =>
Exp lore -> m [ExtType]
expExtType (Apply Name
_ [(SubExp, Diet)]
_ [RetType lore]
rt (Safety, SrcLoc, [SrcLoc])
_) = [ExtType] -> m [ExtType]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExtType] -> m [ExtType]) -> [ExtType] -> m [ExtType]
forall a b. (a -> b) -> a -> b
$ (RetType lore -> ExtType) -> [RetType lore] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape Uniqueness -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl (TypeBase ExtShape Uniqueness -> ExtType)
-> (RetType lore -> TypeBase ExtShape Uniqueness)
-> RetType lore
-> ExtType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetType lore -> TypeBase ExtShape Uniqueness
forall t. DeclExtTyped t => t -> TypeBase ExtShape Uniqueness
declExtTypeOf) [RetType lore]
rt
expExtType (If SubExp
_ BodyT lore
_ BodyT lore
_ IfDec (BranchType lore)
rt) = [ExtType] -> m [ExtType]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExtType] -> m [ExtType]) -> [ExtType] -> m [ExtType]
forall a b. (a -> b) -> a -> b
$ (BranchType lore -> ExtType) -> [BranchType lore] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map BranchType lore -> ExtType
forall t. ExtTyped t => t -> ExtType
extTypeOf ([BranchType lore] -> [ExtType]) -> [BranchType lore] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ IfDec (BranchType lore) -> [BranchType lore]
forall rt. IfDec rt -> [rt]
ifReturns IfDec (BranchType lore)
rt
expExtType (DoLoop [(FParam lore, SubExp)]
ctxmerge [(FParam lore, SubExp)]
valmerge LoopForm lore
_ BodyT lore
_) =
[ExtType] -> m [ExtType]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExtType] -> m [ExtType]) -> [ExtType] -> m [ExtType]
forall a b. (a -> b) -> a -> b
$ [Ident] -> [Ident] -> [ExtType]
loopExtType (((FParam lore, SubExp) -> Ident)
-> [(FParam lore, SubExp)] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map (FParam lore -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent (FParam lore -> Ident)
-> ((FParam lore, SubExp) -> FParam lore)
-> (FParam lore, SubExp)
-> Ident
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam lore, SubExp) -> FParam lore
forall a b. (a, b) -> a
fst) [(FParam lore, SubExp)]
ctxmerge) (((FParam lore, SubExp) -> Ident)
-> [(FParam lore, SubExp)] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map (FParam lore -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent (FParam lore -> Ident)
-> ((FParam lore, SubExp) -> FParam lore)
-> (FParam lore, SubExp)
-> Ident
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam lore, SubExp) -> FParam lore
forall a b. (a, b) -> a
fst) [(FParam lore, SubExp)]
valmerge)
expExtType (BasicOp BasicOp
op) = [Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes ([Type] -> [ExtType]) -> m [Type] -> m [ExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BasicOp -> m [Type]
forall lore (m :: * -> *). HasScope lore m => BasicOp -> m [Type]
primOpType BasicOp
op
expExtType (Op Op lore
op) = Op lore -> m [ExtType]
forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType Op lore
op
expExtTypeSize ::
(Decorations lore, TypedOp (Op lore)) =>
Exp lore ->
Int
expExtTypeSize :: forall lore.
(Decorations lore, TypedOp (Op lore)) =>
Exp lore -> Int
expExtTypeSize = [ExtType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([ExtType] -> Int) -> (Exp lore -> [ExtType]) -> Exp lore -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FeelBad lore [ExtType] -> [ExtType]
forall lore a. FeelBad lore a -> a
feelBad (FeelBad lore [ExtType] -> [ExtType])
-> (Exp lore -> FeelBad lore [ExtType]) -> Exp lore -> [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp lore -> FeelBad lore [ExtType]
forall lore (m :: * -> *).
(HasScope lore m, TypedOp (Op lore)) =>
Exp lore -> m [ExtType]
expExtType
newtype FeelBad lore a = FeelBad {forall lore a. FeelBad lore a -> a
feelBad :: a}
instance Functor (FeelBad lore) where
fmap :: forall a b. (a -> b) -> FeelBad lore a -> FeelBad lore b
fmap a -> b
f = b -> FeelBad lore b
forall lore a. a -> FeelBad lore a
FeelBad (b -> FeelBad lore b)
-> (FeelBad lore a -> b) -> FeelBad lore a -> FeelBad lore b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f (a -> b) -> (FeelBad lore a -> a) -> FeelBad lore a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FeelBad lore a -> a
forall lore a. FeelBad lore a -> a
feelBad
instance Applicative (FeelBad lore) where
pure :: forall a. a -> FeelBad lore a
pure = a -> FeelBad lore a
forall lore a. a -> FeelBad lore a
FeelBad
FeelBad lore (a -> b)
f <*> :: forall a b.
FeelBad lore (a -> b) -> FeelBad lore a -> FeelBad lore b
<*> FeelBad lore a
x = b -> FeelBad lore b
forall lore a. a -> FeelBad lore a
FeelBad (b -> FeelBad lore b) -> b -> FeelBad lore b
forall a b. (a -> b) -> a -> b
$ FeelBad lore (a -> b) -> a -> b
forall lore a. FeelBad lore a -> a
feelBad FeelBad lore (a -> b)
f (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ FeelBad lore a -> a
forall lore a. FeelBad lore a -> a
feelBad FeelBad lore a
x
instance Decorations lore => HasScope lore (FeelBad lore) where
lookupType :: VName -> FeelBad lore Type
lookupType = FeelBad lore Type -> VName -> FeelBad lore Type
forall a b. a -> b -> a
const (FeelBad lore Type -> VName -> FeelBad lore Type)
-> FeelBad lore Type -> VName -> FeelBad lore Type
forall a b. (a -> b) -> a -> b
$ Type -> FeelBad lore Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> FeelBad lore Type) -> Type -> FeelBad lore Type
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int64
askScope :: FeelBad lore (Scope lore)
askScope = Scope lore -> FeelBad lore (Scope lore)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Scope lore
forall a. Monoid a => a
mempty
loopExtType :: [Ident] -> [Ident] -> [ExtType]
loopExtType :: [Ident] -> [Ident] -> [ExtType]
loopExtType [Ident]
ctx [Ident]
val =
[VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes [VName]
inaccessible ([ExtType] -> [ExtType]) -> [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ [Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes ([Type] -> [ExtType]) -> [Type] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ (Ident -> Type) -> [Ident] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> Type
identType [Ident]
val
where
inaccessible :: [VName]
inaccessible = (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
ctx
class TypedOp op where
opType :: HasScope t m => op -> m [ExtType]
instance TypedOp () where
opType :: forall t (m :: * -> *). HasScope t m => () -> m [ExtType]
opType () = [ExtType] -> m [ExtType]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []