module Futhark.Internalise.Defunctionalise (transformProg) where
import Control.Monad
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.List (partition, sortOn)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.IR.Pretty ()
import Futhark.MonadFreshNames
import Futhark.Util (mapAccumLM, nubOrd)
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types (Subst (..), applySubst)
data StaticVal
= Dynamic ParamType
|
LambdaSV (Pat ParamType) ResRetType Exp Env
| RecordSV [(Name, StaticVal)]
|
SumSV Name [StaticVal] [(Name, [ParamType])]
|
DynamicFun (Exp, StaticVal) StaticVal
| IntrinsicSV
| HoleSV StructType SrcLoc
deriving (Int -> StaticVal -> ShowS
[StaticVal] -> ShowS
StaticVal -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [StaticVal] -> ShowS
$cshowList :: [StaticVal] -> ShowS
show :: StaticVal -> [Char]
$cshow :: StaticVal -> [Char]
showsPrec :: Int -> StaticVal -> ShowS
$cshowsPrec :: Int -> StaticVal -> ShowS
Show)
data Binding = Binding
{
Binding -> Maybe ([VName], StructType)
bindingType :: Maybe ([VName], StructType),
Binding -> StaticVal
bindingSV :: StaticVal
}
deriving (Int -> Binding -> ShowS
[Binding] -> ShowS
Binding -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Binding] -> ShowS
$cshowList :: [Binding] -> ShowS
show :: Binding -> [Char]
$cshow :: Binding -> [Char]
showsPrec :: Int -> Binding -> ShowS
$cshowsPrec :: Int -> Binding -> ShowS
Show)
type Env = M.Map VName Binding
localEnv :: Env -> DefM a -> DefM a
localEnv :: forall a. Env -> DefM a -> DefM a
localEnv Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Env
env <>)
localNewEnv :: Env -> DefM a -> DefM a
localNewEnv :: forall a. Env -> DefM a -> DefM a
localNewEnv Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \(Set VName
globals, Env
old_env) ->
(Set VName
globals, forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\VName
k Binding
_ -> VName
k forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals) Env
old_env forall a. Semigroup a => a -> a -> a
<> Env
env)
askEnv :: DefM Env
askEnv :: DefM Env
askEnv = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> b
snd
areGlobal :: [VName] -> DefM a -> DefM a
areGlobal :: forall a. [VName] -> DefM a -> DefM a
areGlobal [VName]
vs = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. Ord a => [a] -> Set a
S.fromList [VName]
vs <>)
replaceTypeSizes ::
M.Map VName SizeSubst ->
TypeBase Size als ->
TypeBase Size als
replaceTypeSizes :: forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Exp -> Exp
onDim
where
onDim :: Exp -> Exp
onDim (Var QualName VName
v Info StructType
typ SrcLoc
loc) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') -> forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' Info StructType
typ SrcLoc
loc
Just (SubstConst Int64
d) -> Integer -> SrcLoc -> Exp
sizeFromInteger (forall a. Integral a => a -> Integer
toInteger Int64
d) SrcLoc
loc
Maybe SizeSubst
Nothing -> forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v Info StructType
typ SrcLoc
loc
onDim Exp
d = Exp
d
replaceStaticValSizes ::
S.Set VName ->
M.Map VName SizeSubst ->
StaticVal ->
StaticVal
replaceStaticValSizes :: Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv =
case StaticVal
sv of
StaticVal
_ | forall k a. Map k a -> Bool
M.null Map VName SizeSubst
orig_substs -> StaticVal
sv
LambdaSV Pat ParamType
param (RetType [VName]
t_dims TypeBase Exp Uniqueness
t) Exp
e Env
closure_env ->
let substs :: Map VName SizeSubst
substs =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Ord k => k -> Map k a -> Map k a
M.delete) Map VName SizeSubst
orig_substs forall a b. (a -> b) -> a -> b
$
forall a. Ord a => [a] -> Set a
S.fromList (forall k a. Map k a -> [k]
M.keys Env
closure_env)
in Pat ParamType -> ResRetType -> Exp -> Env -> StaticVal
LambdaSV
(forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs) Pat ParamType
param)
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
t_dims (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs TypeBase Exp Uniqueness
t))
(Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e)
(forall {k}.
Ord k =>
Map VName SizeSubst -> Map k Binding -> Map k Binding
onEnv Map VName SizeSubst
orig_substs Env
closure_env)
Dynamic ParamType
t ->
ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
orig_substs ParamType
t
RecordSV [(Name, StaticVal)]
fs ->
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs)) [(Name, StaticVal)]
fs
SumSV Name
c [StaticVal]
svs [(Name, [ParamType])]
ts ->
Name -> [StaticVal] -> [(Name, [ParamType])] -> StaticVal
SumSV Name
c (forall a b. (a -> b) -> [a] -> [b]
map (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs) [StaticVal]
svs) forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a -> b) -> a -> b
$ forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
orig_substs) [(Name, [ParamType])]
ts
DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2 ->
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
orig_substs Exp
e, Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv1) forall a b. (a -> b) -> a -> b
$
Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv2
StaticVal
IntrinsicSV ->
StaticVal
IntrinsicSV
HoleSV StructType
t SrcLoc
loc ->
StructType -> SrcLoc -> StaticVal
HoleSV StructType
t SrcLoc
loc
where
tv :: Map VName SizeSubst -> ASTMapper Identity
tv Map VName SizeSubst
substs =
ASTMapper
{ mapOnStructType :: StructType -> Identity StructType
mapOnStructType = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs,
mapOnParamType :: ParamType -> Identity ParamType
mapOnParamType = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs,
mapOnResRetType :: ResRetType -> Identity ResRetType
mapOnResRetType = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnExp :: Exp -> Identity Exp
mapOnExp = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs,
mapOnName :: VName -> Identity VName
mapOnName = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs
}
onName :: Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs VName
v =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') -> forall vn. QualName vn -> vn
qualLeaf QualName VName
v'
Maybe SizeSubst
_ -> VName
v
onExp :: Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs (Var QualName VName
v Info StructType
t SrcLoc
loc) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') ->
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' Info StructType
t SrcLoc
loc
Just (SubstConst Int64
d) ->
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (Int64 -> IntValue
Int64Value (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
d))) SrcLoc
loc
Maybe SizeSubst
Nothing ->
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info StructType
t) SrcLoc
loc
onExp Map VName SizeSubst
substs (Coerce Exp
e TypeExp Info VName
te Info StructType
t SrcLoc
loc) =
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp f vn -> f StructType -> SrcLoc -> ExpBase f vn
Coerce (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e) TypeExp Info VName
te (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info StructType
t) SrcLoc
loc
onExp Map VName SizeSubst
substs (Lambda [Pat ParamType]
params Exp
e Maybe (TypeExp Info VName)
ret (Info (RetType [VName]
t_dims TypeBase Exp Uniqueness
t)) SrcLoc
loc) =
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda
(forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs) [Pat ParamType]
params)
(Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e)
Maybe (TypeExp Info VName)
ret
(forall a. a -> Info a
Info (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
t_dims (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs TypeBase Exp Uniqueness
t)))
SrcLoc
loc
onExp Map VName SizeSubst
substs Exp
e = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (Map VName SizeSubst -> ASTMapper Identity
tv Map VName SizeSubst
substs) Exp
e
onEnv :: Map VName SizeSubst -> Map k Binding -> Map k Binding
onEnv Map VName SizeSubst
substs =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Map VName SizeSubst -> Binding -> Binding
onBinding Map VName SizeSubst
substs))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList
onBinding :: Map VName SizeSubst -> Binding -> Binding
onBinding Map VName SizeSubst
substs (Binding Maybe ([VName], StructType)
t StaticVal
bsv) =
Maybe ([VName], StructType) -> StaticVal -> Binding
Binding
(forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
substs) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ([VName], StructType)
t)
(Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
substs StaticVal
bsv)
restrictEnvTo :: FV -> DefM Env
restrictEnvTo :: FV -> DefM Env
restrictEnvTo FV
fv = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Set VName, Env) -> Env
restrict
where
restrict :: (Set VName, Env) -> Env
restrict (Set VName
globals, Env
env) = forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey VName -> Binding -> Maybe Binding
keep Env
env
where
keep :: VName -> Binding -> Maybe Binding
keep VName
k (Binding Maybe ([VName], StructType)
t StaticVal
sv) = do
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (VName
k forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals) Bool -> Bool -> Bool
&& forall a. Ord a => a -> Set a -> Bool
S.member VName
k (FV -> Set VName
fvVars FV
fv)
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t forall a b. (a -> b) -> a -> b
$ StaticVal -> StaticVal
restrict' StaticVal
sv
restrict' :: StaticVal -> StaticVal
restrict' (Dynamic ParamType
t) =
ParamType -> StaticVal
Dynamic ParamType
t
restrict' (LambdaSV Pat ParamType
pat ResRetType
t Exp
e Env
env) =
Pat ParamType -> ResRetType -> Exp -> Env -> StaticVal
LambdaSV Pat ParamType
pat ResRetType
t Exp
e forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map Binding -> Binding
restrict'' Env
env
restrict' (RecordSV [(Name, StaticVal)]
fields) =
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> StaticVal
restrict') [(Name, StaticVal)]
fields
restrict' (SumSV Name
c [StaticVal]
svs [(Name, [ParamType])]
fields) =
Name -> [StaticVal] -> [(Name, [ParamType])] -> StaticVal
SumSV Name
c (forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> StaticVal
restrict' [StaticVal]
svs) [(Name, [ParamType])]
fields
restrict' (DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2) =
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp
e, StaticVal -> StaticVal
restrict' StaticVal
sv1) forall a b. (a -> b) -> a -> b
$ StaticVal -> StaticVal
restrict' StaticVal
sv2
restrict' StaticVal
IntrinsicSV = StaticVal
IntrinsicSV
restrict' (HoleSV StructType
t SrcLoc
loc) = StructType -> SrcLoc -> StaticVal
HoleSV StructType
t SrcLoc
loc
restrict'' :: Binding -> Binding
restrict'' (Binding Maybe ([VName], StructType)
t StaticVal
sv) = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t forall a b. (a -> b) -> a -> b
$ StaticVal -> StaticVal
restrict' StaticVal
sv
newtype DefM a
= DefM (ReaderT (S.Set VName, Env) (State ([ValBind], VNameSource)) a)
deriving
( forall a b. a -> DefM b -> DefM a
forall a b. (a -> b) -> DefM a -> DefM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> DefM b -> DefM a
$c<$ :: forall a b. a -> DefM b -> DefM a
fmap :: forall a b. (a -> b) -> DefM a -> DefM b
$cfmap :: forall a b. (a -> b) -> DefM a -> DefM b
Functor,
Functor DefM
forall a. a -> DefM a
forall a b. DefM a -> DefM b -> DefM a
forall a b. DefM a -> DefM b -> DefM b
forall a b. DefM (a -> b) -> DefM a -> DefM b
forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. DefM a -> DefM b -> DefM a
$c<* :: forall a b. DefM a -> DefM b -> DefM a
*> :: forall a b. DefM a -> DefM b -> DefM b
$c*> :: forall a b. DefM a -> DefM b -> DefM b
liftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
$cliftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
$c<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
pure :: forall a. a -> DefM a
$cpure :: forall a. a -> DefM a
Applicative,
Applicative DefM
forall a. a -> DefM a
forall a b. DefM a -> DefM b -> DefM b
forall a b. DefM a -> (a -> DefM b) -> DefM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> DefM a
$creturn :: forall a. a -> DefM a
>> :: forall a b. DefM a -> DefM b -> DefM b
$c>> :: forall a b. DefM a -> DefM b -> DefM b
>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
$c>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
Monad,
MonadReader (S.Set VName, Env),
MonadState ([ValBind], VNameSource)
)
instance MonadFreshNames DefM where
putNameSource :: VNameSource -> DefM ()
putNameSource VNameSource
src = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \([ValBind]
x, VNameSource
_) -> ([ValBind]
x, VNameSource
src)
getNameSource :: DefM VNameSource
getNameSource = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a, b) -> b
snd
runDefM :: VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM :: forall a. VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM VNameSource
src (DefM ReaderT (Set VName, Env) (State ([ValBind], VNameSource)) a
m) =
let (a
x, ([ValBind]
vbs, VNameSource
src')) = forall s a. State s a -> s -> (a, s)
runState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Set VName, Env) (State ([ValBind], VNameSource)) a
m forall a. Monoid a => a
mempty) (forall a. Monoid a => a
mempty, VNameSource
src)
in (a
x, VNameSource
src', forall a. [a] -> [a]
reverse [ValBind]
vbs)
addValBind :: ValBind -> DefM ()
addValBind :: ValBind -> DefM ()
addValBind ValBind
vb = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (ValBind
vb :)
liftValDec :: VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()
liftValDec :: VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()
liftValDec VName
fname (RetType [VName]
ret_dims TypeBase Exp Uniqueness
ret) [VName]
dims [Pat ParamType]
pats Exp
body = ValBind -> DefM ()
addValBind ValBind
dec
where
dims' :: [TypeParamBase VName]
dims' = forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) [VName]
dims
bound_here :: Set VName
bound_here = forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ [VName]
dims forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall t. Pat t -> [VName]
patNames [Pat ParamType]
pats
mkExt :: VName -> Maybe VName
mkExt VName
v
| Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ VName
v forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound_here = forall a. a -> Maybe a
Just VName
v
mkExt VName
_ = forall a. Maybe a
Nothing
rettype_st :: ResRetType
rettype_st = forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe VName
mkExt (forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars forall a b. (a -> b) -> a -> b
$ forall u. TypeBase Exp u -> FV
freeInType TypeBase Exp Uniqueness
ret) forall a. [a] -> [a] -> [a]
++ [VName]
ret_dims) TypeBase Exp Uniqueness
ret
dec :: ValBind
dec =
ValBind
{ valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = forall a. Maybe a
Nothing,
valBindName :: VName
valBindName = VName
fname,
valBindRetDecl :: Maybe (TypeExp Info VName)
valBindRetDecl = forall a. Maybe a
Nothing,
valBindRetType :: Info ResRetType
valBindRetType = forall a. a -> Info a
Info ResRetType
rettype_st,
valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
dims',
valBindParams :: [Pat ParamType]
valBindParams = [Pat ParamType]
pats,
valBindBody :: Exp
valBindBody = Exp
body,
valBindDoc :: Maybe DocComment
valBindDoc = forall a. Maybe a
Nothing,
valBindAttrs :: [AttrInfo VName]
valBindAttrs = forall a. Monoid a => a
mempty,
valBindLocation :: SrcLoc
valBindLocation = forall a. Monoid a => a
mempty
}
lookupVar :: StructType -> VName -> DefM StaticVal
lookupVar :: StructType -> VName -> DefM StaticVal
lookupVar StructType
t VName
x = do
Env
env <- DefM Env
askEnv
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Env
env of
Just (Binding (Just ([VName]
dims, StructType
sv_t)) StaticVal
sv) -> do
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [VName] -> StructType -> StructType -> StaticVal -> m StaticVal
instStaticVal Set VName
globals [VName]
dims StructType
t StructType
sv_t StaticVal
sv
Just (Binding Maybe ([VName], StructType)
Nothing StaticVal
sv) ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure StaticVal
sv
Maybe Binding
Nothing
| VName -> Int
baseTag VName
x forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag -> forall (f :: * -> *) a. Applicative f => a -> f a
pure StaticVal
IntrinsicSV
| Bool
otherwise ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
arraySizes :: StructType -> S.Set VName
arraySizes :: StructType -> Set VName
arraySizes (Scalar Arrow {}) = forall a. Monoid a => a
mempty
arraySizes (Scalar (Record Map Name StructType
fields)) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
arraySizes Map Name StructType
fields
arraySizes (Scalar (Sum Map Name [StructType]
cs)) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
arraySizes) Map Name [StructType]
cs
arraySizes (Scalar (TypeVar NoUniqueness
_ QualName VName
_ [TypeArg Exp]
targs)) =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map TypeArg Exp -> Set VName
f [TypeArg Exp]
targs
where
f :: TypeArg Exp -> Set VName
f (TypeArgDim (Var QualName VName
d Info StructType
_ SrcLoc
_)) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
d
f TypeArgDim {} = forall a. Monoid a => a
mempty
f (TypeArgType StructType
t) = StructType -> Set VName
arraySizes StructType
t
arraySizes (Scalar Prim {}) = forall a. Monoid a => a
mempty
arraySizes (Array NoUniqueness
_ Shape Exp
shape ScalarTypeBase Exp NoUniqueness
t) =
StructType -> Set VName
arraySizes (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase Exp NoUniqueness
t) forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> Set VName
dimName (forall dim. Shape dim -> [dim]
shapeDims Shape Exp
shape)
where
dimName :: Size -> S.Set VName
dimName :: Exp -> Set VName
dimName (Var QualName VName
qn Info StructType
_ SrcLoc
_) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
dimName Exp
_ = forall a. Monoid a => a
mempty
patternArraySizes :: Pat ParamType -> S.Set VName
patternArraySizes :: Pat ParamType -> Set VName
patternArraySizes = StructType -> Set VName
arraySizes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. Pat (TypeBase Exp u) -> StructType
patternStructType
data SizeSubst
= SubstNamed (QualName VName)
| SubstConst Int64
deriving (SizeSubst -> SizeSubst -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SizeSubst -> SizeSubst -> Bool
$c/= :: SizeSubst -> SizeSubst -> Bool
== :: SizeSubst -> SizeSubst -> Bool
$c== :: SizeSubst -> SizeSubst -> Bool
Eq, Eq SizeSubst
SizeSubst -> SizeSubst -> Bool
SizeSubst -> SizeSubst -> Ordering
SizeSubst -> SizeSubst -> SizeSubst
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
min :: SizeSubst -> SizeSubst -> SizeSubst
$cmin :: SizeSubst -> SizeSubst -> SizeSubst
max :: SizeSubst -> SizeSubst -> SizeSubst
$cmax :: SizeSubst -> SizeSubst -> SizeSubst
>= :: SizeSubst -> SizeSubst -> Bool
$c>= :: SizeSubst -> SizeSubst -> Bool
> :: SizeSubst -> SizeSubst -> Bool
$c> :: SizeSubst -> SizeSubst -> Bool
<= :: SizeSubst -> SizeSubst -> Bool
$c<= :: SizeSubst -> SizeSubst -> Bool
< :: SizeSubst -> SizeSubst -> Bool
$c< :: SizeSubst -> SizeSubst -> Bool
compare :: SizeSubst -> SizeSubst -> Ordering
$ccompare :: SizeSubst -> SizeSubst -> Ordering
Ord, Int -> SizeSubst -> ShowS
[SizeSubst] -> ShowS
SizeSubst -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [SizeSubst] -> ShowS
$cshowList :: [SizeSubst] -> ShowS
show :: SizeSubst -> [Char]
$cshow :: SizeSubst -> [Char]
showsPrec :: Int -> SizeSubst -> ShowS
$cshowsPrec :: Int -> SizeSubst -> ShowS
Show)
dimMapping ::
Monoid a =>
TypeBase Size a ->
TypeBase Size a ->
M.Map VName SizeSubst
dimMapping :: forall a.
Monoid a =>
TypeBase Exp a -> TypeBase Exp a -> Map VName SizeSubst
dimMapping TypeBase Exp a
t1 TypeBase Exp a
t2 = forall s a. State s a -> s -> s
execState (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 {t :: * -> *} {f :: * -> *} {k} {f :: * -> *} {f :: * -> *}.
(Foldable t, MonadState (Map k SizeSubst) f, Ord k) =>
t VName -> ExpBase f k -> ExpBase f VName -> f (ExpBase f k)
f TypeBase Exp a
t1 TypeBase Exp a
t2) forall a. Monoid a => a
mempty
where
f :: t VName -> ExpBase f k -> ExpBase f VName -> f (ExpBase f k)
f t VName
bound ExpBase f k
d1 (Var QualName VName
d2 f StructType
_ SrcLoc
_)
| forall vn. QualName vn -> vn
qualLeaf QualName VName
d2 forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
bound = forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f k
d1
f t VName
_ (Var QualName k
d1 f StructType
typ SrcLoc
loc) (Var QualName VName
d2 f StructType
_ SrcLoc
_) = do
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall vn. QualName vn -> vn
qualLeaf QualName k
d1) forall a b. (a -> b) -> a -> b
$ QualName VName -> SizeSubst
SubstNamed QualName VName
d2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName k
d1 f StructType
typ SrcLoc
loc
f t VName
_ (Var QualName k
d1 f StructType
typ SrcLoc
loc) (IntLit Integer
d2 f StructType
_ SrcLoc
_) = do
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall vn. QualName vn -> vn
qualLeaf QualName k
d1) forall a b. (a -> b) -> a -> b
$ Int64 -> SizeSubst
SubstConst forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger Integer
d2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName k
d1 f StructType
typ SrcLoc
loc
f t VName
_ ExpBase f k
d ExpBase f VName
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f k
d
dimMapping' ::
Monoid a =>
TypeBase Size a ->
TypeBase Size a ->
M.Map VName VName
dimMapping' :: forall a.
Monoid a =>
TypeBase Exp a -> TypeBase Exp a -> Map VName VName
dimMapping' TypeBase Exp a
t1 TypeBase Exp a
t2 = forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe SizeSubst -> Maybe VName
f forall a b. (a -> b) -> a -> b
$ forall a.
Monoid a =>
TypeBase Exp a -> TypeBase Exp a -> Map VName SizeSubst
dimMapping TypeBase Exp a
t1 TypeBase Exp a
t2
where
f :: SizeSubst -> Maybe VName
f (SubstNamed QualName VName
d) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
d
f SizeSubst
_ = forall a. Maybe a
Nothing
sizesToRename :: StaticVal -> S.Set VName
sizesToRename :: StaticVal -> Set VName
sizesToRename (DynamicFun (Exp
_, StaticVal
sv1) StaticVal
sv2) =
StaticVal -> Set VName
sizesToRename StaticVal
sv1 forall a. Semigroup a => a -> a -> a
<> StaticVal -> Set VName
sizesToRename StaticVal
sv2
sizesToRename StaticVal
IntrinsicSV =
forall a. Monoid a => a
mempty
sizesToRename HoleSV {} =
forall a. Monoid a => a
mempty
sizesToRename Dynamic {} =
forall a. Monoid a => a
mempty
sizesToRename (RecordSV [(Name, StaticVal)]
fs) =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (StaticVal -> Set VName
sizesToRename forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, StaticVal)]
fs
sizesToRename (SumSV Name
_ [StaticVal]
svs [(Name, [ParamType])]
_) =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StaticVal -> Set VName
sizesToRename [StaticVal]
svs
sizesToRename (LambdaSV Pat ParamType
param ResRetType
_ Exp
_ Env
_) =
FV -> Set VName
fvVars (forall u. Pat (TypeBase Exp u) -> FV
freeInPat Pat ParamType
param)
combineTypeShapes ::
(Monoid as) =>
TypeBase Size as ->
TypeBase Size as ->
TypeBase Size as
combineTypeShapes :: forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes (Scalar (Record Map Name (TypeBase Exp as)
ts1)) (Scalar (Record Map Name (TypeBase Exp as)
ts2))
| forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp as)
ts1 forall a. Eq a => a -> a -> Bool
== forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp as)
ts2 =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record forall a b. (a -> b) -> a -> b
$
forall a b k. (a -> b) -> Map k a -> Map k b
M.map
(forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes)
(forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name (TypeBase Exp as)
ts1 Map Name (TypeBase Exp as)
ts2)
combineTypeShapes (Scalar (Sum Map Name [TypeBase Exp as]
cs1)) (Scalar (Sum Map Name [TypeBase Exp as]
cs2))
| forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Exp as]
cs1 forall a. Eq a => a -> a -> Bool
== forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Exp as]
cs2 =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum forall a b. (a -> b) -> a -> b
$
forall a b k. (a -> b) -> Map k a -> Map k b
M.map
(forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes)
(forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name [TypeBase Exp as]
cs1 Map Name [TypeBase Exp as]
cs2)
combineTypeShapes (Scalar (Arrow as
als1 PName
p1 Diet
d1 StructType
a1 (RetType [VName]
dims1 TypeBase Exp Uniqueness
b1))) (Scalar (Arrow as
als2 PName
_p2 Diet
_d2 StructType
a2 (RetType [VName]
_ TypeBase Exp Uniqueness
b2))) =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow
(as
als1 forall a. Semigroup a => a -> a -> a
<> as
als2)
PName
p1
Diet
d1
(forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes StructType
a1 StructType
a2)
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims1 (forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes TypeBase Exp Uniqueness
b1 TypeBase Exp Uniqueness
b2))
combineTypeShapes (Scalar (TypeVar as
u QualName VName
v [TypeArg Exp]
targs1)) (Scalar (TypeVar as
_ QualName VName
_ [TypeArg Exp]
targs2)) =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar as
u QualName VName
v forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeArg Exp -> TypeArg Exp -> TypeArg Exp
f [TypeArg Exp]
targs1 [TypeArg Exp]
targs2
where
f :: TypeArg Exp -> TypeArg Exp -> TypeArg Exp
f (TypeArgType StructType
t1) (TypeArgType StructType
t2) = forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes StructType
t1 StructType
t2)
f TypeArg Exp
targ TypeArg Exp
_ = TypeArg Exp
targ
combineTypeShapes (Array as
u Shape Exp
shape1 ScalarTypeBase Exp NoUniqueness
et1) (Array as
_ Shape Exp
_shape2 ScalarTypeBase Exp NoUniqueness
et2) =
forall u dim. u -> Shape dim -> TypeBase dim u -> TypeBase dim u
arrayOfWithAliases
as
u
Shape Exp
shape1
(forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes (forall dim u1 u2. TypeBase dim u1 -> u2 -> TypeBase dim u2
setUniqueness (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase Exp NoUniqueness
et1) as
u) (forall dim u1 u2. TypeBase dim u1 -> u2 -> TypeBase dim u2
setUniqueness (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase Exp NoUniqueness
et2) as
u))
combineTypeShapes TypeBase Exp as
_ TypeBase Exp as
t = TypeBase Exp as
t
instStaticVal ::
MonadFreshNames m =>
S.Set VName ->
[VName] ->
StructType ->
StructType ->
StaticVal ->
m StaticVal
instStaticVal :: forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [VName] -> StructType -> StructType -> StaticVal -> m StaticVal
instStaticVal Set VName
globals [VName]
dims StructType
t StructType
sv_t StaticVal
sv = do
Map VName SizeSubst
fresh_substs <-
forall {f :: * -> *}.
MonadFreshNames f =>
[VName] -> f (Map VName SizeSubst)
mkSubsts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
globals) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$
forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims forall a. Semigroup a => a -> a -> a
<> StaticVal -> Set VName
sizesToRename StaticVal
sv
let dims' :: [VName]
dims' = forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
fresh_substs) [VName]
dims
isDim :: VName -> p -> Bool
isDim VName
k p
_ = VName
k forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
dims'
dim_substs :: Map VName SizeSubst
dim_substs =
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey forall {p}. VName -> p -> Bool
isDim forall a b. (a -> b) -> a -> b
$ forall a.
Monoid a =>
TypeBase Exp a -> TypeBase Exp a -> Map VName SizeSubst
dimMapping (forall als.
Map VName SizeSubst -> TypeBase Exp als -> TypeBase Exp als
replaceTypeSizes Map VName SizeSubst
fresh_substs StructType
sv_t) StructType
t
replace :: SizeSubst -> SizeSubst
replace (SubstNamed QualName VName
k) = forall a. a -> Maybe a -> a
fromMaybe (QualName VName -> SizeSubst
SubstNamed QualName VName
k) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
k) Map VName SizeSubst
dim_substs
replace SizeSubst
k = SizeSubst
k
substs :: Map VName SizeSubst
substs = forall a b k. (a -> b) -> Map k a -> Map k b
M.map SizeSubst -> SizeSubst
replace Map VName SizeSubst
fresh_substs forall a. Semigroup a => a -> a -> a
<> Map VName SizeSubst
dim_substs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
substs StaticVal
sv
where
mkSubsts :: [VName] -> f (Map VName SizeSubst)
mkSubsts [VName]
names =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (QualName VName -> SizeSubst
SubstNamed forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName [VName]
names
onName :: Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs VName
v =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') -> forall vn. QualName vn -> vn
qualLeaf QualName VName
v'
Maybe SizeSubst
_ -> VName
v
defuncFun ::
[VName] ->
[Pat ParamType] ->
Exp ->
ResRetType ->
SrcLoc ->
DefM (Exp, StaticVal)
defuncFun :: [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
tparams [Pat ParamType]
pats Exp
e0 ResRetType
ret SrcLoc
loc = do
let (Pat ParamType
pat, ResRetType
ret', Exp
e0') = case [Pat ParamType]
pats of
[] -> forall a. HasCallStack => [Char] -> a
error [Char]
"Received a lambda with no parameters."
[Pat ParamType
pat'] -> (Pat ParamType
pat', ResRetType
ret, Exp
e0)
(Pat ParamType
pat' : [Pat ParamType]
pats') ->
( Pat ParamType
pat',
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a b. a -> b -> a
const Uniqueness
Nonunique) forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> ResRetType -> StructType
funType [Pat ParamType]
pats' ResRetType
ret,
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats' Exp
e0 forall a. Maybe a
Nothing (forall a. a -> Info a
Info ResRetType
ret) SrcLoc
loc
)
let used :: FV
used =
Exp -> FV
freeInExp (forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats Exp
e0 forall a. Maybe a
Nothing (forall a. a -> Info a
Info ResRetType
ret) SrcLoc
loc)
FV -> Set VName -> FV
`freeWithout` forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams
Env
used_env <- FV -> DefM Env
restrictEnvTo FV
used
let sizes_of_arrays :: Set VName
sizes_of_arrays =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (StructType -> Set VName
arraySizes forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> StructType
structTypeFromSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV) Env
used_env
forall a. Semigroup a => a -> a -> a
<> Pat ParamType -> Set VName
patternArraySizes Pat ParamType
pat
notSize :: VName -> Bool
notSize = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
sizes_of_arrays)
([FieldBase Info VName]
fields, Env
env) =
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [(a, b)] -> ([a], [b])
unzip
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Bool
notSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)
forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Env
used_env
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fields SrcLoc
loc,
Pat ParamType -> ResRetType -> Exp -> Env -> StaticVal
LambdaSV Pat ParamType
pat ResRetType
ret' Exp
e0' Env
env
)
where
closureFromDynamicFun :: (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ (DynamicFun (Exp
clsr_env, StaticVal
sv) StaticVal
_)) =
let name :: Name
name = [Char] -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString VName
vn
in ( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name Exp
clsr_env forall a. Monoid a => a
mempty,
(VName
vn, Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing StaticVal
sv)
)
closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
let name :: Name
name = [Char] -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString VName
vn
tp' :: StructType
tp' = StaticVal -> StructType
structTypeFromSV StaticVal
sv
in ( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
Name
name
(forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
vn) (forall a. a -> Info a
Info StructType
tp') forall a. Monoid a => a
mempty)
forall a. Monoid a => a
mempty,
(VName
vn, Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing StaticVal
sv)
)
defuncExp :: Exp -> DefM (Exp, StaticVal)
defuncExp :: Exp -> DefM (Exp, StaticVal)
defuncExp e :: Exp
e@Literal {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp e :: Exp
e@IntLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp e :: Exp
e@FloatLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp e :: Exp
e@StringLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp (Parens Exp
e SrcLoc
loc) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens Exp
e' SrcLoc
loc, StaticVal
sv)
defuncExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn Exp
e' SrcLoc
loc, StaticVal
sv)
defuncExp (TupLit [Exp]
es SrcLoc
loc) = do
([Exp]
es', [StaticVal]
svs) <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [Exp]
es' SrcLoc
loc, [(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [StaticVal]
svs)
defuncExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
([FieldBase Info VName]
fs', [(Name, StaticVal)]
names_svs) <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField [FieldBase Info VName]
fs
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fs' SrcLoc
loc, [(Name, StaticVal)] -> StaticVal
RecordSV [(Name, StaticVal)]
names_svs)
where
defuncField :: FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField (RecordFieldExplicit Name
vn Exp
e SrcLoc
loc') = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
vn Exp
e' SrcLoc
loc', (Name
vn, StaticVal
sv))
defuncField (RecordFieldImplicit VName
vn (Info StructType
t) SrcLoc
loc') = do
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t) VName
vn
case StaticVal
sv of
DynamicFun (Exp
e, StaticVal
sv') StaticVal
_ ->
let vn' :: Name
vn' = VName -> Name
baseName VName
vn
in forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
vn' Exp
e SrcLoc
loc',
(Name
vn', StaticVal
sv')
)
StaticVal
_ ->
let tp :: Info StructType
tp = forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> StructType
structTypeFromSV StaticVal
sv
in forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
vn -> f StructType -> SrcLoc -> FieldBase f vn
RecordFieldImplicit VName
vn Info StructType
tp SrcLoc
loc', (VName -> Name
baseName VName
vn, StaticVal
sv))
defuncExp (ArrayLit [Exp]
es t :: Info StructType
t@(Info StructType
t') SrcLoc
loc) = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp' [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
ArrayLit [Exp]
es' Info StructType
t SrcLoc
loc, ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe StructType
t')
defuncExp (AppExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl SrcLoc
loc) Info AppRes
res) = do
Exp
e1' <- Exp -> DefM Exp
defuncExp' Exp
e1
Maybe Exp
me' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp' Maybe Exp
me
Inclusiveness Exp
incl' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp' Inclusiveness Exp
incl
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Exp
e1' Maybe Exp
me' Inclusiveness Exp
incl' SrcLoc
loc) Info AppRes
res,
ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ AppRes -> StructType
appResType forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo Info AppRes
res
)
defuncExp e :: Exp
e@(Var QualName VName
qn (Info StructType
t) SrcLoc
loc) = do
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t) (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
case StaticVal
sv of
DynamicFun {} -> do
([Pat ParamType]
params, Exp
body, ResRetType
ret) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique StructType
t) Exp
e
[VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [Pat ParamType]
params Exp
body ResRetType
ret forall a. Monoid a => a
mempty
StaticVal
IntrinsicSV -> do
([Pat ParamType]
pats, Exp
body, ResRetType
tp) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique StructType
t) Exp
e
Exp -> DefM (Exp, StaticVal)
defuncExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info ResRetType
tp) forall a. Monoid a => a
mempty
HoleSV StructType
_ SrcLoc
hole_loc ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (forall a. a -> Info a
Info StructType
t) SrcLoc
hole_loc, StaticVal
sv)
StaticVal
_ ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info (StaticVal -> StructType
structTypeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncExp (Hole (Info StructType
t) SrcLoc
loc) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (forall a. a -> Info a
Info StructType
t) SrcLoc
loc, StructType -> SrcLoc -> StaticVal
HoleSV StructType
t SrcLoc
loc)
defuncExp (Ascript Exp
e0 TypeExp Info VName
tydecl SrcLoc
loc)
| forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> StructType
typeOf Exp
e0) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp f vn -> SrcLoc -> ExpBase f vn
Ascript Exp
e0' TypeExp Info VName
tydecl SrcLoc
loc, StaticVal
sv)
| Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (Coerce Exp
e0 TypeExp Info VName
tydecl Info StructType
t SrcLoc
loc)
| forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> StructType
typeOf Exp
e0) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp f vn -> f StructType -> SrcLoc -> ExpBase f vn
Coerce Exp
e0' TypeExp Info VName
tydecl Info StructType
t SrcLoc
loc, StaticVal
sv)
| Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e1 Exp
e2 SrcLoc
loc) (Info (AppRes StructType
t [VName]
retext))) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
let env :: Env
env = Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
pat) StaticVal
sv1
pat' :: Pat ParamType
pat' = Pat ParamType -> StaticVal -> Pat ParamType
updatePat (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
pat) StaticVal
sv1
(Exp
e2', StaticVal
sv2) <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
let mapping :: Map VName VName
mapping = forall a.
Monoid a =>
TypeBase Exp a -> TypeBase Exp a -> Map VName VName
dimMapping' (Exp -> StructType
typeOf Exp
e2) StructType
t
subst :: VName -> Maybe (Subst t)
subst VName
v = forall t. Exp -> Subst t
ExpSubst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
mapping
t' :: StructType
t' = forall a. Substitutable a => TypeSubs -> a -> a
applySubst forall {t}. VName -> Maybe (Subst t)
subst forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e2'
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct Pat ParamType
pat') Exp
e1' Exp
e2' SrcLoc
loc) (forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes StructType
t' [VName]
retext)), StaticVal
sv2)
defuncExp (AppExp (LetFun VName
vn ([TypeParamBase VName], [Pat ParamType],
Maybe (TypeExp Info VName), Info ResRetType, Exp)
_ Exp
_ SrcLoc
_) Info AppRes
_) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"defuncExp: Unexpected LetFun: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show VName
vn
defuncExp (AppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) Info AppRes
res) = do
(Exp
e1', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
(Exp
e2', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
(Exp
e3', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e3
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
e1' Exp
e2' Exp
e3' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
defuncExp (AppExp (Apply Exp
f NonEmpty (Info (Diet, Maybe VName), Exp)
args SrcLoc
loc) (Info AppRes
appres)) =
Exp
-> NonEmpty ((Diet, Maybe VName), Exp)
-> AppRes
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncApply Exp
f (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. Info a -> a
unInfo) NonEmpty (Info (Diet, Maybe VName), Exp)
args) AppRes
appres SrcLoc
loc
defuncExp (Negate Exp
e0 SrcLoc
loc) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate Exp
e0' SrcLoc
loc, StaticVal
sv)
defuncExp (Not Exp
e0 SrcLoc
loc) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not Exp
e0' SrcLoc
loc, StaticVal
sv)
defuncExp (Lambda [Pat ParamType]
pats Exp
e0 Maybe (TypeExp Info VName)
_ (Info ResRetType
ret) SrcLoc
loc) =
[VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [Pat ParamType]
pats Exp
e0 ResRetType
ret SrcLoc
loc
defuncExp OpSection {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp OpSectionLeft {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp OpSectionRight {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp ProjectSection {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected projection section."
defuncExp IndexSection {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected projection section."
defuncExp (AppExp (DoLoop [VName]
sparams Pat ParamType
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 SrcLoc
loc) Info AppRes
res) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
let env1 :: Env
env1 = Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Pat ParamType
pat StaticVal
sv1
(LoopFormBase Info VName
form', Env
env2) <- case LoopFormBase Info VName
form of
For IdentBase Info VName StructType
v Exp
e2 -> do
Exp
e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
IdentBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName StructType
v Exp
e2', forall {k} {u}. IdentBase Info k (TypeBase Exp u) -> Map k Binding
envFromIdent IdentBase Info VName StructType
v)
ForIn PatBase Info VName StructType
pat2 Exp
e2 -> do
Exp
e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
ForIn PatBase Info VName StructType
pat2 Exp
e2', Pat ParamType -> Env
envFromPat forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
pat2)
While Exp
e2 -> do
Exp
e2' <- forall a. Env -> DefM a -> DefM a
localEnv Env
env1 forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While Exp
e2', forall a. Monoid a => a
mempty)
(Exp
e3', StaticVal
sv) <- forall a. Env -> DefM a -> DefM a
localEnv (Env
env1 forall a. Semigroup a => a -> a -> a
<> Env
env2) forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e3
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[VName]
-> PatBase f vn ParamType
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop [VName]
sparams Pat ParamType
pat Exp
e1' LoopFormBase Info VName
form' Exp
e3' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
where
envFromIdent :: IdentBase Info k (TypeBase Exp u) -> Map k Binding
envFromIdent (Ident k
vn (Info TypeBase Exp u
tp) SrcLoc
_) =
forall k a. k -> a -> Map k a
M.singleton k
vn forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe TypeBase Exp u
tp
defuncExp e :: Exp
e@(AppExp BinOp {} Info AppRes
_) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"defuncExp: unexpected binary operator: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Exp
e
defuncExp (Project Name
vn Exp
e0 tp :: Info StructType
tp@(Info StructType
tp') SrcLoc
loc) = do
(Exp
e0', StaticVal
sv0) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
case StaticVal
sv0 of
RecordSV [(Name, StaticVal)]
svs -> case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
vn [(Name, StaticVal)]
svs of
Just StaticVal
sv -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> StructType
structTypeFromSV StaticVal
sv) SrcLoc
loc, StaticVal
sv)
Maybe StaticVal
Nothing -> forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
Dynamic ParamType
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' Info StructType
tp SrcLoc
loc, ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe StructType
tp')
HoleSV StructType
_ SrcLoc
hloc -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' Info StructType
tp SrcLoc
loc, StructType -> SrcLoc -> StaticVal
HoleSV StructType
tp' SrcLoc
hloc)
StaticVal
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Projection of an expression with static value " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv0
defuncExp (AppExp (LetWith IdentBase Info VName StructType
id1 IdentBase Info VName StructType
id2 SliceBase Info VName
idxs Exp
e1 Exp
body SrcLoc
loc) Info AppRes
res) = do
Exp
e1' <- Exp -> DefM Exp
defuncExp' Exp
e1
SliceBase Info VName
idxs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs
let id1_binding :: Binding
id1_binding =
Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo forall a b. (a -> b) -> a -> b
$ forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType IdentBase Info VName StructType
id1
(Exp
body', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localEnv (forall k a. k -> a -> Map k a
M.singleton (forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase Info VName StructType
id1) Binding
id1_binding) forall a b. (a -> b) -> a -> b
$
Exp -> DefM (Exp, StaticVal)
defuncExp Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
IdentBase f vn StructType
-> IdentBase f vn StructType
-> SliceBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetWith IdentBase Info VName StructType
id1 IdentBase Info VName StructType
id2 SliceBase Info VName
idxs' Exp
e1' Exp
body' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
defuncExp expr :: Exp
expr@(AppExp (Index Exp
e0 SliceBase Info VName
idxs SrcLoc
loc) Info AppRes
res) = do
Exp
e0' <- Exp -> DefM Exp
defuncExp' Exp
e0
SliceBase Info VName
idxs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index Exp
e0' SliceBase Info VName
idxs' SrcLoc
loc) Info AppRes
res,
ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
expr
)
defuncExp (Update Exp
e1 SliceBase Info VName
idxs Exp
e2 SrcLoc
loc) = do
(Exp
e1', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
SliceBase Info VName
idxs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs
Exp
e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update Exp
e1' SliceBase Info VName
idxs' Exp
e2' SrcLoc
loc, StaticVal
sv)
defuncExp (RecordUpdate Exp
e1 [Name]
fs Exp
e2 Info StructType
_ SrcLoc
loc) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
(Exp
e2', StaticVal
sv2) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
let sv :: StaticVal
sv = StaticVal -> StaticVal -> [Name] -> StaticVal
staticField StaticVal
sv1 StaticVal
sv2 [Name]
fs
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
RecordUpdate Exp
e1' [Name]
fs Exp
e2' (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> StructType
structTypeFromSV StaticVal
sv1) SrcLoc
loc,
StaticVal
sv
)
where
staticField :: StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (RecordSV [(Name, StaticVal)]
svs) StaticVal
sv2 (Name
f : [Name]
fs') =
case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
f [(Name, StaticVal)]
svs of
Just StaticVal
sv ->
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$
(Name
f, StaticVal -> StaticVal -> [Name] -> StaticVal
staticField StaticVal
sv StaticVal
sv2 [Name]
fs') forall a. a -> [a] -> [a]
: forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Eq a => a -> a -> Bool
/= Name
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Name, StaticVal)]
svs
Maybe StaticVal
Nothing -> forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
staticField (Dynamic t :: ParamType
t@(Scalar Record {})) StaticVal
sv2 fs' :: [Name]
fs'@(Name
_ : [Name]
_) =
StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (ParamType -> StaticVal
svFromType ParamType
t) StaticVal
sv2 [Name]
fs'
staticField StaticVal
_ StaticVal
sv2 [Name]
_ = StaticVal
sv2
defuncExp (Assert Exp
e1 Exp
e2 Info Text
desc SrcLoc
loc) = do
(Exp
e1', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
(Exp
e2', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert Exp
e1' Exp
e2' Info Text
desc SrcLoc
loc, StaticVal
sv)
defuncExp (Constr Name
name [Exp]
es (Info sum_t :: StructType
sum_t@(Scalar (Sum Map Name [StructType]
all_fs))) SrcLoc
loc) = do
([Exp]
es', [StaticVal]
svs) <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
let sv :: StaticVal
sv =
Name -> [StaticVal] -> [(Name, [ParamType])] -> StaticVal
SumSV Name
name [StaticVal]
svs forall a b. (a -> b) -> a -> b
$
forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$
Name
name forall k a. Ord k => k -> Map k a -> Map k a
`M.delete` forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall als. Monoid als => TypeBase Exp als -> TypeBase Exp als
defuncType)) Map Name [StructType]
all_fs
sum_t' :: StructType
sum_t' = forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes StructType
sum_t (StaticVal -> StructType
structTypeFromSV StaticVal
sv)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
Constr Name
name [Exp]
es' (forall a. a -> Info a
Info StructType
sum_t') SrcLoc
loc, StaticVal
sv)
where
defuncType ::
Monoid als =>
TypeBase Size als ->
TypeBase Size als
defuncType :: forall als. Monoid als => TypeBase Exp als -> TypeBase Exp als
defuncType (Array als
u Shape Exp
shape ScalarTypeBase Exp NoUniqueness
t) = forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array als
u Shape Exp
shape (forall als.
Monoid als =>
ScalarTypeBase Exp als -> ScalarTypeBase Exp als
defuncScalar ScalarTypeBase Exp NoUniqueness
t)
defuncType (Scalar ScalarTypeBase Exp als
t) = forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall als.
Monoid als =>
ScalarTypeBase Exp als -> ScalarTypeBase Exp als
defuncScalar ScalarTypeBase Exp als
t
defuncScalar ::
Monoid als =>
ScalarTypeBase Size als ->
ScalarTypeBase Size als
defuncScalar :: forall als.
Monoid als =>
ScalarTypeBase Exp als -> ScalarTypeBase Exp als
defuncScalar (Record Map Name (TypeBase Exp als)
fs) = forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map forall als. Monoid als => TypeBase Exp als -> TypeBase Exp als
defuncType Map Name (TypeBase Exp als)
fs
defuncScalar Arrow {} = forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record forall a. Monoid a => a
mempty
defuncScalar (Sum Map Name [TypeBase Exp als]
fs) = forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall als. Monoid als => TypeBase Exp als -> TypeBase Exp als
defuncType) Map Name [TypeBase Exp als]
fs
defuncScalar (Prim PrimType
t) = forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t
defuncScalar (TypeVar als
u QualName VName
tn [TypeArg Exp]
targs) = forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar als
u QualName VName
tn [TypeArg Exp]
targs
defuncExp (Constr Name
name [Exp]
_ (Info StructType
t) SrcLoc
loc) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Constructor "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
name
forall a. [a] -> [a] -> [a]
++ [Char]
" given type "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString StructType
t
forall a. [a] -> [a] -> [a]
++ [Char]
" at "
forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr SrcLoc
loc
defuncExp (AppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
res) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
let bad :: a
bad = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"No case matches StaticVal\n" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show StaticVal
sv
NonEmpty (CaseBase Info VName, StaticVal)
csPairs <-
forall a. a -> Maybe a -> a
fromMaybe forall {a}. a
bad forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StaticVal
-> CaseBase Info VName
-> DefM (Maybe (CaseBase Info VName, StaticVal))
defuncCase StaticVal
sv) (forall a. NonEmpty a -> [a]
NE.toList NonEmpty (CaseBase Info VName)
cs)
let cs' :: NonEmpty (CaseBase Info VName)
cs' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst NonEmpty (CaseBase Info VName, StaticVal)
csPairs
sv' :: StaticVal
sv' = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> a
NE.head NonEmpty (CaseBase Info VName, StaticVal)
csPairs
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match Exp
e' NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc) Info AppRes
res, StaticVal
sv')
defuncExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
info Exp
e' SrcLoc
loc, StaticVal
sv)
defuncExp' :: Exp -> DefM Exp
defuncExp' :: Exp -> DefM Exp
defuncExp' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> DefM (Exp, StaticVal)
defuncExp
defuncCase :: StaticVal -> Case -> DefM (Maybe (Case, StaticVal))
defuncCase :: StaticVal
-> CaseBase Info VName
-> DefM (Maybe (CaseBase Info VName, StaticVal))
defuncCase StaticVal
sv (CasePat PatBase Info VName StructType
p Exp
e SrcLoc
loc) = do
let p' :: Pat ParamType
p' = Pat ParamType -> StaticVal -> Pat ParamType
updatePat (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
p) StaticVal
sv
case Pat ParamType -> StaticVal -> Maybe Env
matchPatSV (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
p) StaticVal
sv of
Just Env
env -> do
(Exp
e', StaticVal
sv') <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct Pat ParamType
p') Exp
e' SrcLoc
loc, StaticVal
sv')
Maybe Env
Nothing ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
defuncSoacExp :: Exp -> DefM Exp
defuncSoacExp :: Exp -> DefM Exp
defuncSoacExp e :: Exp
e@OpSection {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@OpSectionLeft {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@OpSectionRight {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@ProjectSection {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp (Parens Exp
e SrcLoc
loc) =
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> DefM Exp
defuncSoacExp Exp
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
defuncSoacExp (Lambda [Pat ParamType]
params Exp
e0 Maybe (TypeExp Info VName)
decl Info ResRetType
tp SrcLoc
loc) = do
let env :: Env
env = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> Env
envFromPat [Pat ParamType]
params
Exp
e0' <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncSoacExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
params Exp
e0' Maybe (TypeExp Info VName)
decl Info ResRetType
tp SrcLoc
loc
defuncSoacExp Exp
e
| Scalar Arrow {} <- Exp -> StructType
typeOf Exp
e = do
([Pat ParamType]
pats, Exp
body, ResRetType
tp) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e) Exp
e
let env :: Env
env = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> Env
envFromPat [Pat ParamType]
pats
Exp
body' <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncExp' Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats Exp
body' forall a. Maybe a
Nothing (forall a. a -> Info a
Info ResRetType
tp) forall a. Monoid a => a
mempty
| Bool
otherwise = Exp -> DefM Exp
defuncExp' Exp
e
etaExpand :: ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand :: ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ResRetType
e_t Exp
e = do
let ([(PName, (Diet, StructType))]
ps, ResRetType
ret) = forall {dim}.
RetTypeBase dim Uniqueness
-> ([(PName, (Diet, TypeBase dim NoUniqueness))],
RetTypeBase dim Uniqueness)
getType ResRetType
e_t
([VName]
_, ([Pat ParamType]
params, [Exp]
vars)) <- forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall a b. [(a, b)] -> ([a], [b])
unzip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (t :: * -> *) acc x y.
(Monad m, Traversable t) =>
(acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
mapAccumLM forall {m :: * -> *} {u} {b}.
MonadFreshNames m =>
[VName]
-> (PName, (u, TypeBase Exp b))
-> m ([VName], (PatBase Info VName (TypeBase Exp u), Exp))
f [] [(PName, (Diet, StructType))]
ps
[VName]
ext' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName forall a b. (a -> b) -> a -> b
$ forall dim as. RetTypeBase dim as -> [VName]
retDims ResRetType
ret
let extsubst :: Map VName (Subst t)
extsubst =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip (forall dim as. RetTypeBase dim as -> [VName]
retDims ResRetType
ret) forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall t. Exp -> Subst t
ExpSubst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName) [VName]
ext'
ret' :: ResRetType
ret' = forall a. Substitutable a => TypeSubs -> a -> a
applySubst (forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` forall {t}. Map VName (Subst t)
extsubst) ResRetType
ret
e' :: Exp
e' =
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
Exp
e
(forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(PName, (Diet, StructType))]
ps) (forall a. a -> [a]
repeat forall a. Maybe a
Nothing) [Exp]
vars)
(StructType -> [VName] -> AppRes
AppRes (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall a b. (a -> b) -> a -> b
$ forall dim as. RetTypeBase dim as -> TypeBase dim as
retType ResRetType
ret') [VName]
ext')
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Pat ParamType]
params, Exp
e', ResRetType
ret)
where
getType :: RetTypeBase dim Uniqueness
-> ([(PName, (Diet, TypeBase dim NoUniqueness))],
RetTypeBase dim Uniqueness)
getType (RetType [VName]
_ (Scalar (Arrow Uniqueness
_ PName
p Diet
d TypeBase dim NoUniqueness
t1 RetTypeBase dim Uniqueness
t2))) =
let ([(PName, (Diet, TypeBase dim NoUniqueness))]
ps, RetTypeBase dim Uniqueness
r) = RetTypeBase dim Uniqueness
-> ([(PName, (Diet, TypeBase dim NoUniqueness))],
RetTypeBase dim Uniqueness)
getType RetTypeBase dim Uniqueness
t2
in ((PName
p, (Diet
d, TypeBase dim NoUniqueness
t1)) forall a. a -> [a] -> [a]
: [(PName, (Diet, TypeBase dim NoUniqueness))]
ps, RetTypeBase dim Uniqueness
r)
getType RetTypeBase dim Uniqueness
t = ([], RetTypeBase dim Uniqueness
t)
f :: [VName]
-> (PName, (u, TypeBase Exp b))
-> m ([VName], (PatBase Info VName (TypeBase Exp u), Exp))
f [VName]
prev (PName
p, (u
d, TypeBase Exp b
t)) = do
let t' :: TypeBase Exp u
t' = forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a b. a -> b -> a
const u
d) TypeBase Exp b
t
VName
x <- case PName
p of
Named VName
x | VName
x forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
prev -> forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
x
PName
_ -> forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"eta_p"
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( VName
x forall a. a -> [a] -> [a]
: [VName]
prev,
( forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
x (forall a. a -> Info a
Info TypeBase Exp u
t') forall a. Monoid a => a
mempty,
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
x) (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp u
t') forall a. Monoid a => a
mempty
)
)
defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex (DimFix Exp
e1) = forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
defuncDimIndex (DimSlice Maybe Exp
me1 Maybe Exp
me2 Maybe Exp
me3) =
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me2 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me3
where
defunc' :: Maybe Exp -> DefM (Maybe Exp)
defunc' = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp'
defuncLet ::
[VName] ->
[Pat ParamType] ->
Exp ->
ResRetType ->
DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
defuncLet :: [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> DefM
([VName], [Pat ParamType], Exp, StaticVal, TypeBase Exp Uniqueness)
defuncLet [VName]
dims ps :: [Pat ParamType]
ps@(Pat ParamType
pat : [Pat ParamType]
pats) Exp
body (RetType [VName]
ret_dims TypeBase Exp Uniqueness
rettype)
| forall d u. Pat (TypeBase d u) -> Bool
patternOrderZero Pat ParamType
pat = do
let bound_by_pat :: VName -> Bool
bound_by_pat = (forall a. Ord a => a -> Set a -> Bool
`S.member` FV -> Set VName
fvVars (forall u. Pat (TypeBase Exp u) -> FV
freeInPat Pat ParamType
pat))
([VName]
pat_dims, [VName]
rest_dims) = forall a. (a -> Bool) -> [a] -> ([a], [a])
partition VName -> Bool
bound_by_pat [VName]
dims
env :: Env
env = Pat ParamType -> Env
envFromPat Pat ParamType
pat forall a. Semigroup a => a -> a -> a
<> [VName] -> Env
envFromDimNames [VName]
pat_dims
([VName]
rest_dims', [Pat ParamType]
pats', Exp
body', StaticVal
sv, TypeBase Exp Uniqueness
sv_t) <-
forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> DefM
([VName], [Pat ParamType], Exp, StaticVal, TypeBase Exp Uniqueness)
defuncLet [VName]
rest_dims [Pat ParamType]
pats Exp
body forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims TypeBase Exp Uniqueness
rettype
(Exp, StaticVal)
closure <- [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [Pat ParamType]
ps Exp
body (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims TypeBase Exp Uniqueness
rettype) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( [VName]
pat_dims forall a. [a] -> [a] -> [a]
++ [VName]
rest_dims',
Pat ParamType
pat forall a. a -> [a] -> [a]
: [Pat ParamType]
pats',
Exp
body',
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp, StaticVal)
closure StaticVal
sv,
TypeBase Exp Uniqueness
sv_t
)
| Bool
otherwise = do
(Exp
e, StaticVal
sv) <- [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [Pat ParamType]
ps Exp
body (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims TypeBase Exp Uniqueness
rettype) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], Exp
e, StaticVal
sv, StaticVal -> TypeBase Exp Uniqueness
resTypeFromSV StaticVal
sv)
defuncLet [VName]
_ [] Exp
body (RetType [VName]
_ TypeBase Exp Uniqueness
rettype) = do
(Exp
body', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( [],
[],
Exp
body',
StaticVal -> ParamType -> StaticVal
imposeType StaticVal
sv forall a b. (a -> b) -> a -> b
$ TypeBase Exp Uniqueness -> ParamType
resToParam TypeBase Exp Uniqueness
rettype,
StaticVal -> TypeBase Exp Uniqueness
resTypeFromSV StaticVal
sv
)
where
imposeType :: StaticVal -> ParamType -> StaticVal
imposeType Dynamic {} ParamType
t =
ParamType -> StaticVal
Dynamic ParamType
t
imposeType (RecordSV [(Name, StaticVal)]
fs1) (Scalar (Record Map Name ParamType
fs2)) =
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith StaticVal -> ParamType -> StaticVal
imposeType (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, StaticVal)]
fs1) Map Name ParamType
fs2
imposeType StaticVal
sv ParamType
_ = StaticVal
sv
instAnySizes :: MonadFreshNames m => [Pat ParamType] -> m [Pat ParamType]
instAnySizes :: forall (m :: * -> *).
MonadFreshNames m =>
[Pat ParamType] -> m [Pat ParamType]
instAnySizes = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a b. (a -> b) -> a -> b
$ 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 forall {m :: * -> *}. MonadFreshNames m => Exp -> m Exp
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure
where
onDim :: Exp -> m Exp
onDim Exp
d
| Exp
d forall a. Eq a => a -> a -> Bool
== Exp
anySize = do
VName
v <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"size"
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (forall v. v -> QualName v
qualName VName
v) forall a. Monoid a => a
mempty
onDim Exp
d = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d
unboundSizes :: S.Set VName -> [Pat ParamType] -> [VName]
unboundSizes :: Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
params = forall a. Ord a => [a] -> [a]
nubOrd forall a b. (a -> b) -> a -> b
$ forall s a. State s a -> s -> s
execState (forall {f :: * -> *} {d}.
[PatBase Info VName (TypeBase (ExpBase f VName) d)]
-> StateT
[VName]
Identity
[PatBase Info VName (TypeBase (ExpBase f VName) d)]
f [Pat ParamType]
params) []
where
f :: [PatBase Info VName (TypeBase (ExpBase f VName) d)]
-> StateT
[VName]
Identity
[PatBase Info VName (TypeBase (ExpBase f VName) d)]
f = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a b. (a -> b) -> a -> b
$ 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 forall {m :: * -> *} {f :: * -> *}.
MonadState [VName] m =>
ExpBase f VName -> m (ExpBase f VName)
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure
bound :: Set VName
bound = Set VName
bound_sizes forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)
onDim :: ExpBase f VName -> m (ExpBase f VName)
onDim (Var QualName VName
d f StructType
typ SrcLoc
loc) = do
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall vn. QualName vn -> vn
qualLeaf QualName VName
d forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (forall vn. QualName vn -> vn
qualLeaf QualName VName
d :)
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
d f StructType
typ SrcLoc
loc
onDim ExpBase f VName
d = forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f VName
d
unRetType :: ResRetType -> DefM AppRes
unRetType :: ResRetType -> DefM AppRes
unRetType (RetType [] TypeBase Exp Uniqueness
t) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
t) []
unRetType (RetType [VName]
ext TypeBase Exp Uniqueness
t) = do
[VName]
ext' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName [VName]
ext
let extsubst :: Map VName (Subst t)
extsubst =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
ext forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall t. Exp -> Subst t
ExpSubst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName) [VName]
ext'
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (forall a. Substitutable a => TypeSubs -> a -> a
applySubst (forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` forall {t}. Map VName (Subst t)
extsubst) forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
t) [VName]
ext'
defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal)
defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal)
defuncApplyFunction e :: Exp
e@(Var QualName VName
qn (Info StructType
t) SrcLoc
loc) Int
num_args = do
let ([ParamType]
argtypes, StructType
rettype) = forall dim as.
TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType StructType
t
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t) (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
case StaticVal
sv of
DynamicFun (Exp, StaticVal)
_ StaticVal
_
| StaticVal -> Int -> Bool
fullyApplied StaticVal
sv Int
num_args -> do
let ([ParamType]
argtypes', TypeBase Exp Uniqueness
rettype') = StaticVal -> [ParamType] -> ([ParamType], TypeBase Exp Uniqueness)
dynamicFunType StaticVal
sv [ParamType]
argtypes
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info ([ParamType] -> ResRetType -> StructType
foldFunType [ParamType]
argtypes' forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
rettype')) SrcLoc
loc, StaticVal
sv)
| forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall dim as. TypeBase dim as -> Bool
orderZero [ParamType]
argtypes,
forall dim as. TypeBase dim as -> Bool
orderZero StructType
rettype -> do
([Pat ParamType]
params, Exp
body, ResRetType
ret) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique StructType
t) Exp
e
[VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [Pat ParamType]
params Exp
body ResRetType
ret forall a. Monoid a => a
mempty
| Bool
otherwise -> do
VName
fname <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName forall a b. (a -> b) -> a -> b
$ [Char]
"dyn_" forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
let ([Pat ParamType]
pats, Exp
e0, StaticVal
sv') = [Char] -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun (forall a. Pretty a => a -> [Char]
prettyString QualName VName
qn) StaticVal
sv Int
num_args
([ParamType]
argtypes', TypeBase Exp Uniqueness
rettype') = StaticVal -> [ParamType] -> ([ParamType], TypeBase Exp Uniqueness)
dynamicFunType StaticVal
sv' [ParamType]
argtypes
dims' :: [VName]
dims' = forall a. Monoid a => a
mempty
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let bound_sizes :: Set VName
bound_sizes = forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims' forall a. Semigroup a => a -> a -> a
<> Set VName
globals
[Pat ParamType]
pats' <- forall (m :: * -> *).
MonadFreshNames m =>
[Pat ParamType] -> m [Pat ParamType]
instAnySizes [Pat ParamType]
pats
VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()
liftValDec VName
fname (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
rettype') ([VName]
dims' forall a. [a] -> [a] -> [a]
++ Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
pats') [Pat ParamType]
pats' Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var
(forall v. v -> QualName v
qualName VName
fname)
(forall a. a -> Info a
Info ([ParamType] -> ResRetType -> StructType
foldFunType [ParamType]
argtypes' forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
rettype'))
SrcLoc
loc,
StaticVal
sv'
)
StaticVal
IntrinsicSV -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, StaticVal
IntrinsicSV)
StaticVal
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info (StaticVal -> StructType
structTypeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncApplyFunction Exp
e Int
_ = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
liftedName :: Int -> Exp -> String
liftedName :: Int -> Exp -> [Char]
liftedName Int
i (Var QualName VName
f Info StructType
_ SrcLoc
_) =
[Char]
"defunc_" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
i forall a. [a] -> [a] -> [a]
++ [Char]
"_" forall a. [a] -> [a] -> [a]
++ VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
f)
liftedName Int
i (AppExp (Apply Exp
f NonEmpty (Info (Diet, Maybe VName), Exp)
_ SrcLoc
_) Info AppRes
_) =
Int -> Exp -> [Char]
liftedName (Int
i forall a. Num a => a -> a -> a
+ Int
1) Exp
f
liftedName Int
_ Exp
_ = [Char]
"defunc"
defuncApplyArg ::
String ->
(Exp, StaticVal) ->
(((Diet, Maybe VName), Exp), [ParamType]) ->
DefM (Exp, StaticVal)
defuncApplyArg :: [Char]
-> (Exp, StaticVal)
-> (((Diet, Maybe VName), Exp), [ParamType])
-> DefM (Exp, StaticVal)
defuncApplyArg [Char]
fname_s (Exp
f', LambdaSV Pat ParamType
pat ResRetType
lam_e_t Exp
lam_e Env
closure_env) (((Diet
d, Maybe VName
argext), Exp
arg), [ParamType]
_) = do
(Exp
arg', StaticVal
arg_sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
arg
let env' :: Env
env' = Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Pat ParamType
pat StaticVal
arg_sv
dims :: [VName]
dims = forall a. Monoid a => a
mempty
(Exp
lam_e', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localNewEnv (Env
env' forall a. Semigroup a => a -> a -> a
<> Env
closure_env) forall a b. (a -> b) -> a -> b
$
Exp -> DefM (Exp, StaticVal)
defuncExp Exp
lam_e
let closure_pat :: Pat ParamType
closure_pat = [VName] -> Env -> Pat ParamType
buildEnvPat [VName]
dims Env
closure_env
pat' :: Pat ParamType
pat' = Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
arg_sv
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let params :: [Pat ParamType]
params = [Pat ParamType
closure_pat, Pat ParamType
pat']
lifted_rettype :: ResRetType
lifted_rettype =
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (forall dim as. RetTypeBase dim as -> [VName]
retDims ResRetType
lam_e_t) forall a b. (a -> b) -> a -> b
$
forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes (forall dim as. RetTypeBase dim as -> TypeBase dim as
retType ResRetType
lam_e_t) (StaticVal -> TypeBase Exp Uniqueness
resTypeFromSV StaticVal
sv)
already_bound :: Set VName
already_bound =
Set VName
globals forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList ([VName]
dims forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)
more_dims :: [VName]
more_dims =
forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$
forall a. (a -> Bool) -> Set a -> Set a
S.filter (forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
already_bound) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> Set VName
patternArraySizes [Pat ParamType]
params
let bound_sizes :: Set VName
bound_sizes = forall a. Ord a => [a] -> Set a
S.fromList ([VName]
dims forall a. Semigroup a => a -> a -> a
<> [VName]
more_dims) forall a. Semigroup a => a -> a -> a
<> Set VName
globals
[Pat ParamType]
params' <- forall (m :: * -> *).
MonadFreshNames m =>
[Pat ParamType] -> m [Pat ParamType]
instAnySizes [Pat ParamType]
params
VName
fname <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
fname_s
VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()
liftValDec
VName
fname
ResRetType
lifted_rettype
([VName]
dims forall a. [a] -> [a] -> [a]
++ [VName]
more_dims forall a. [a] -> [a] -> [a]
++ Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
params')
[Pat ParamType]
params'
Exp
lam_e'
let f_t :: StructType
f_t = forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
f'
arg_t :: StructType
arg_t = forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
arg'
fname_t :: StructType
fname_t = [ParamType] -> ResRetType -> StructType
foldFunType [forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe StructType
f_t, forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
d StructType
arg_t] ResRetType
lifted_rettype
fname' :: Exp
fname' = forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
fname) (forall a. a -> Info a
Info StructType
fname_t) (forall a. Located a => a -> SrcLoc
srclocOf Exp
arg)
AppRes
callret <- ResRetType -> DefM AppRes
unRetType ResRetType
lifted_rettype
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
fname' [(Diet
Observe, forall a. Maybe a
Nothing, Exp
f'), (Diet
Observe, Maybe VName
argext, Exp
arg')] AppRes
callret,
StaticVal
sv
)
defuncApplyArg [Char]
_ (Exp
f', DynamicFun (Exp, StaticVal)
_ StaticVal
sv) (((Diet
d, Maybe VName
argext), Exp
arg), [ParamType]
argtypes) = do
(Exp
arg', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
arg
let ([ParamType]
argtypes', TypeBase Exp Uniqueness
rettype) = StaticVal -> [ParamType] -> ([ParamType], TypeBase Exp Uniqueness)
dynamicFunType StaticVal
sv [ParamType]
argtypes
restype :: StructType
restype = [ParamType] -> ResRetType -> StructType
foldFunType [ParamType]
argtypes' (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
rettype)
callret :: AppRes
callret = StructType -> [VName] -> AppRes
AppRes StructType
restype []
apply_e :: Exp
apply_e = forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
f' [(Diet
d, Maybe VName
argext, Exp
arg')] AppRes
callret
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
apply_e, StaticVal
sv)
defuncApplyArg [Char]
fname_s (Exp
_, StaticVal
sv) (((Diet, Maybe VName), Exp), [ParamType])
_ =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"defuncApplyArg: cannot apply StaticVal\n"
forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show StaticVal
sv
forall a. Semigroup a => a -> a -> a
<> [Char]
"\nFunction name: "
forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString [Char]
fname_s
updateReturn :: AppRes -> Exp -> Exp
updateReturn :: AppRes -> Exp -> Exp
updateReturn (AppRes StructType
ret1 [VName]
ext1) (AppExp AppExpBase Info VName
apply (Info (AppRes StructType
ret2 [VName]
ext2))) =
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp AppExpBase Info VName
apply forall a b. (a -> b) -> a -> b
$ forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes StructType
ret1 StructType
ret2) ([VName]
ext1 forall a. Semigroup a => a -> a -> a
<> [VName]
ext2)
updateReturn AppRes
_ Exp
e = Exp
e
defuncApply :: Exp -> NE.NonEmpty ((Diet, Maybe VName), Exp) -> AppRes -> SrcLoc -> DefM (Exp, StaticVal)
defuncApply :: Exp
-> NonEmpty ((Diet, Maybe VName), Exp)
-> AppRes
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncApply Exp
f NonEmpty ((Diet, Maybe VName), Exp)
args AppRes
appres SrcLoc
loc = do
(Exp
f', StaticVal
f_sv) <- Exp -> Int -> DefM (Exp, StaticVal)
defuncApplyFunction Exp
f (forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmpty ((Diet, Maybe VName), Exp)
args)
case StaticVal
f_sv of
StaticVal
IntrinsicSV -> do
NonEmpty (Info (Diet, Maybe VName), Exp)
args' <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. a -> Info a
Info) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Exp -> DefM Exp
defuncSoacExp) NonEmpty ((Diet, Maybe VName), Exp)
args
let e' :: Exp
e' = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Diet, Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f' NonEmpty (Info (Diet, Maybe VName), Exp)
args' SrcLoc
loc) (forall a. a -> Info a
Info AppRes
appres)
Exp -> DefM (Exp, StaticVal)
intrinsicOrHole Exp
e'
HoleSV {} -> do
NonEmpty (Info (Diet, Maybe VName), Exp)
args' <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. a -> Info a
Info) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> DefM (Exp, StaticVal)
defuncExp) NonEmpty ((Diet, Maybe VName), Exp)
args
let e' :: Exp
e' = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Diet, Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f' NonEmpty (Info (Diet, Maybe VName), Exp)
args' SrcLoc
loc) (forall a. a -> Info a
Info AppRes
appres)
Exp -> DefM (Exp, StaticVal)
intrinsicOrHole Exp
e'
StaticVal
_ -> do
let fname :: [Char]
fname = Int -> Exp -> [Char]
liftedName Int
0 Exp
f
([ParamType]
argtypes, StructType
_) = forall dim as.
TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
f
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a b. (a -> b) -> a -> b
$ AppRes -> Exp -> Exp
updateReturn AppRes
appres) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Char]
-> (Exp, StaticVal)
-> (((Diet, Maybe VName), Exp), [ParamType])
-> DefM (Exp, StaticVal)
defuncApplyArg [Char]
fname) (Exp
f', StaticVal
f_sv) forall a b. (a -> b) -> a -> b
$
forall a b. NonEmpty a -> NonEmpty b -> NonEmpty (a, b)
NE.zip NonEmpty ((Diet, Maybe VName), Exp)
args forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a. Foldable f => f a -> NonEmpty [a]
NE.tails [ParamType]
argtypes
where
intrinsicOrHole :: Exp -> DefM (Exp, StaticVal)
intrinsicOrHole Exp
e' = do
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall dim as.
TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType forall a b. (a -> b) -> a -> b
$ AppRes -> StructType
appResType AppRes
appres
then forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e', ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe forall a b. (a -> b) -> a -> b
$ AppRes -> StructType
appResType AppRes
appres)
else do
([Pat ParamType]
pats, Exp
body, ResRetType
tp) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e') Exp
e'
Exp -> DefM (Exp, StaticVal)
defuncExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info ResRetType
tp) forall a. Monoid a => a
mempty
fullyApplied :: StaticVal -> Int -> Bool
fullyApplied :: StaticVal -> Int -> Bool
fullyApplied (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) Int
depth
| Int
depth forall a. Eq a => a -> a -> Bool
== Int
0 = Bool
False
| Int
depth forall a. Ord a => a -> a -> Bool
> Int
0 = StaticVal -> Int -> Bool
fullyApplied StaticVal
sv (Int
depth forall a. Num a => a -> a -> a
- Int
1)
fullyApplied StaticVal
_ Int
_ = Bool
True
liftDynFun :: String -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun :: [Char] -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun [Char]
_ (DynamicFun (Exp
e, StaticVal
sv) StaticVal
_) Int
0 = ([], Exp
e, StaticVal
sv)
liftDynFun [Char]
s (DynamicFun clsr :: (Exp, StaticVal)
clsr@(Exp
_, LambdaSV Pat ParamType
pat ResRetType
_ Exp
_ Env
_) StaticVal
sv) Int
d
| Int
d forall a. Ord a => a -> a -> Bool
> Int
0 =
let ([Pat ParamType]
pats, Exp
e', StaticVal
sv') = [Char] -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun [Char]
s StaticVal
sv (Int
d forall a. Num a => a -> a -> a
- Int
1)
in (Pat ParamType
pat forall a. a -> [a] -> [a]
: [Pat ParamType]
pats, Exp
e', (Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp, StaticVal)
clsr StaticVal
sv')
liftDynFun [Char]
s StaticVal
sv Int
d =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
s
forall a. [a] -> [a] -> [a]
++ [Char]
" Tried to lift a StaticVal "
forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
take Int
100 (forall a. Show a => a -> [Char]
show StaticVal
sv)
forall a. [a] -> [a] -> [a]
++ [Char]
", but expected a dynamic function.\n"
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Int
d
envFromPat :: Pat ParamType -> Env
envFromPat :: Pat ParamType -> Env
envFromPat Pat ParamType
pat = case Pat ParamType
pat of
TuplePat [Pat ParamType]
ps SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> Env
envFromPat [Pat ParamType]
ps
RecordPat [(Name, Pat ParamType)]
fs SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Pat ParamType -> Env
envFromPat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, Pat ParamType)]
fs
PatParens Pat ParamType
p SrcLoc
_ -> Pat ParamType -> Env
envFromPat Pat ParamType
p
PatAttr AttrInfo VName
_ Pat ParamType
p SrcLoc
_ -> Pat ParamType -> Env
envFromPat Pat ParamType
p
Id VName
vn (Info ParamType
t) SrcLoc
_ -> forall k a. k -> a -> Map k a
M.singleton VName
vn forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic ParamType
t
Wildcard Info ParamType
_ SrcLoc
_ -> forall a. Monoid a => a
mempty
PatAscription Pat ParamType
p TypeExp Info VName
_ SrcLoc
_ -> Pat ParamType -> Env
envFromPat Pat ParamType
p
PatLit {} -> forall a. Monoid a => a
mempty
PatConstr Name
_ Info ParamType
_ [Pat ParamType]
ps SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> Env
envFromPat [Pat ParamType]
ps
envFromDimNames :: [VName] -> Env
envFromDimNames :: [VName] -> Env
envFromDimNames = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. a -> [a]
repeat Binding
d)
where
d :: Binding
d = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
buildEnvPat :: [VName] -> Env -> Pat ParamType
buildEnvPat :: [VName] -> Env -> Pat ParamType
buildEnvPat [VName]
sizes Env
env = forall (f :: * -> *) vn t.
[(Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat (forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (Name, Pat ParamType)
buildField forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Env
env) forall a. Monoid a => a
mempty
where
buildField :: (VName, Binding) -> (Name, Pat ParamType)
buildField (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
( [Char] -> Name
nameFromString (forall a. Pretty a => a -> [Char]
prettyString VName
vn),
if VName
vn forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
sizes
then forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) forall a. Monoid a => a
mempty
else forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
vn (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) forall a. Monoid a => a
mempty
)
typeFromSV :: StaticVal -> ParamType
typeFromSV :: StaticVal -> ParamType
typeFromSV (Dynamic ParamType
tp) =
ParamType
tp
typeFromSV (LambdaSV Pat ParamType
_ ResRetType
_ Exp
_ Env
env) =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ([Char] -> Name
nameFromString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) (StaticVal -> ParamType
typeFromSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV)) forall a b. (a -> b) -> a -> b
$
forall k a. Map k a -> [(k, a)]
M.toList Env
env
typeFromSV (RecordSV [(Name, StaticVal)]
ls) =
let ts :: [(Name, ParamType)]
ts = forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> ParamType
typeFromSV) [(Name, StaticVal)]
ls
in forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, ParamType)]
ts
typeFromSV (DynamicFun (Exp
_, StaticVal
sv) StaticVal
_) =
StaticVal -> ParamType
typeFromSV StaticVal
sv
typeFromSV (SumSV Name
name [StaticVal]
svs [(Name, [ParamType])]
fields) =
let svs' :: [ParamType]
svs' = forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> ParamType
typeFromSV [StaticVal]
svs
in forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
name [ParamType]
svs' forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [ParamType])]
fields
typeFromSV (HoleSV StructType
t SrcLoc
_) =
forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe StructType
t
typeFromSV StaticVal
IntrinsicSV =
forall a. HasCallStack => [Char] -> a
error [Char]
"Tried to get the type from the static value of an intrinsic."
resTypeFromSV :: StaticVal -> ResType
resTypeFromSV :: StaticVal -> TypeBase Exp Uniqueness
resTypeFromSV = ParamType -> TypeBase Exp Uniqueness
paramToRes forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> ParamType
typeFromSV
structTypeFromSV :: StaticVal -> StructType
structTypeFromSV :: StaticVal -> StructType
structTypeFromSV = forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> ParamType
typeFromSV
paramTypeFromSV :: StaticVal -> ParamType
paramTypeFromSV :: StaticVal -> ParamType
paramTypeFromSV = StaticVal -> ParamType
typeFromSV
dynamicFunType :: StaticVal -> [ParamType] -> ([ParamType], ResType)
dynamicFunType :: StaticVal -> [ParamType] -> ([ParamType], TypeBase Exp Uniqueness)
dynamicFunType (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) (ParamType
p : [ParamType]
ps) =
let ([ParamType]
ps', TypeBase Exp Uniqueness
ret) = StaticVal -> [ParamType] -> ([ParamType], TypeBase Exp Uniqueness)
dynamicFunType StaticVal
sv [ParamType]
ps
in (ParamType
p forall a. a -> [a] -> [a]
: [ParamType]
ps', TypeBase Exp Uniqueness
ret)
dynamicFunType StaticVal
sv [ParamType]
_ = ([], StaticVal -> TypeBase Exp Uniqueness
resTypeFromSV StaticVal
sv)
matchPatSV :: Pat ParamType -> StaticVal -> Maybe Env
matchPatSV :: Pat ParamType -> StaticVal -> Maybe Env
matchPatSV (TuplePat [Pat ParamType]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls) =
forall a. Monoid a => [a] -> a
mconcat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (\Pat ParamType
p (Name
_, StaticVal
sv) -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
p StaticVal
sv) [Pat ParamType]
ps [(Name, StaticVal)]
ls
matchPatSV (RecordPat [(Name, Pat ParamType)]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls)
| [(Name, Pat ParamType)]
ps' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, Pat ParamType)]
ps,
[(Name, StaticVal)]
ls' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, StaticVal)]
ls,
forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Name, Pat ParamType)]
ps' forall a. Eq a => a -> a -> Bool
== forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Name, StaticVal)]
ls' =
forall a. Monoid a => [a] -> a
mconcat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (\(Name
_, Pat ParamType
p) (Name
_, StaticVal
sv) -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
p StaticVal
sv) [(Name, Pat ParamType)]
ps' [(Name, StaticVal)]
ls'
matchPatSV (PatParens Pat ParamType
pat SrcLoc
_) StaticVal
sv = Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
pat StaticVal
sv
matchPatSV (PatAttr AttrInfo VName
_ Pat ParamType
pat SrcLoc
_) StaticVal
sv = Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
pat StaticVal
sv
matchPatSV (Id VName
vn (Info ParamType
t) SrcLoc
_) StaticVal
sv =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
if forall dim as. TypeBase dim as -> Bool
orderZero ParamType
t
then Env
dim_env forall a. Semigroup a => a -> a -> a
<> forall k a. k -> a -> Map k a
M.singleton VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic ParamType
t)
else Env
dim_env forall a. Semigroup a => a -> a -> a
<> forall k a. k -> a -> Map k a
M.singleton VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing StaticVal
sv)
where
dim_env :: Env
dim_env =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (,Binding
i64) forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars forall a b. (a -> b) -> a -> b
$ forall u. TypeBase Exp u -> FV
freeInType ParamType
t
i64 :: Binding
i64 = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
matchPatSV (Wildcard Info ParamType
_ SrcLoc
_) StaticVal
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
matchPatSV (PatAscription Pat ParamType
pat TypeExp Info VName
_ SrcLoc
_) StaticVal
sv = Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
pat StaticVal
sv
matchPatSV PatLit {} StaticVal
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
matchPatSV (PatConstr Name
c1 Info ParamType
_ [Pat ParamType]
ps SrcLoc
_) (SumSV Name
c2 [StaticVal]
ls [(Name, [ParamType])]
fs)
| Name
c1 forall a. Eq a => a -> a -> Bool
== Name
c2 =
forall a. Monoid a => [a] -> a
mconcat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Pat ParamType -> StaticVal -> Maybe Env
matchPatSV [Pat ParamType]
ps [StaticVal]
ls
| Just [ParamType]
_ <- forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
c1 [(Name, [ParamType])]
fs =
forall a. Maybe a
Nothing
| Bool
otherwise =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"matchPatSV: missing constructor in type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
c1
matchPatSV (PatConstr Name
c1 Info ParamType
_ [Pat ParamType]
ps SrcLoc
_) (Dynamic (Scalar (Sum Map Name [ParamType]
fs)))
| Just [ParamType]
ts <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c1 Map Name [ParamType]
fs =
if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall dim as. TypeBase dim as -> Bool
orderZero [ParamType]
ts
then forall a. Monoid a => [a] -> a
mconcat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Pat ParamType -> StaticVal -> Maybe Env
matchPatSV [Pat ParamType]
ps (forall a b. (a -> b) -> [a] -> [b]
map ParamType -> StaticVal
svFromType [ParamType]
ts)
else forall a. Maybe a
Nothing
| Bool
otherwise =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"matchPatSV: missing constructor in type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
c1
matchPatSV Pat ParamType
pat (Dynamic ParamType
t) = Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
pat forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
svFromType ParamType
t
matchPatSV Pat ParamType
pat (HoleSV StructType
t SrcLoc
_) = Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
pat forall a b. (a -> b) -> a -> b
$ ParamType -> StaticVal
svFromType forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe StructType
t
matchPatSV Pat ParamType
pat StaticVal
sv =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Tried to match pattern\n"
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Pat ParamType
pat
forall a. [a] -> [a] -> [a]
++ [Char]
"\n with static value\n"
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv
alwaysMatchPatSV :: Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV :: Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Pat ParamType
pat StaticVal
sv = forall a. a -> Maybe a -> a
fromMaybe forall {a}. a
bad forall a b. (a -> b) -> a -> b
$ Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Pat ParamType
pat StaticVal
sv
where
bad :: a
bad = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unlines [forall a. Pretty a => a -> [Char]
prettyString Pat ParamType
pat, [Char]
"cannot match StaticVal", forall a. Show a => a -> [Char]
show StaticVal
sv]
updatePat :: Pat ParamType -> StaticVal -> Pat ParamType
updatePat :: Pat ParamType -> StaticVal -> Pat ParamType
updatePat (TuplePat [Pat ParamType]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs) =
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Pat ParamType -> StaticVal -> Pat ParamType
updatePat [Pat ParamType]
ps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, StaticVal)]
svs) SrcLoc
loc
updatePat (RecordPat [(Name, Pat ParamType)]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs)
| [(Name, Pat ParamType)]
ps' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, Pat ParamType)]
ps,
[(Name, StaticVal)]
svs' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, StaticVal)]
svs =
forall (f :: * -> *) vn t.
[(Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat
(forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
n, Pat ParamType
p) (Name
_, StaticVal
sv) -> (Name
n, Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
p StaticVal
sv)) [(Name, Pat ParamType)]
ps' [(Name, StaticVal)]
svs')
SrcLoc
loc
updatePat (PatParens Pat ParamType
pat SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn t.
PatBase f vn t -> SrcLoc -> PatBase f vn t
PatParens (Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
sv) SrcLoc
loc
updatePat (PatAttr AttrInfo VName
attr Pat ParamType
pat SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn t.
AttrInfo vn -> PatBase f vn t -> SrcLoc -> PatBase f vn t
PatAttr AttrInfo VName
attr (Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
sv) SrcLoc
loc
updatePat (Id VName
vn (Info ParamType
tp) SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
vn (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall {dim} {u}.
TypeBase dim u -> TypeBase dim u -> TypeBase dim u
comb ParamType
tp forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) SrcLoc
loc
where
comb :: TypeBase dim u -> TypeBase dim u -> TypeBase dim u
comb (Scalar Arrow {}) TypeBase dim u
t2 = TypeBase dim u
t2
comb (Scalar (Record Map Name (TypeBase dim u)
m1)) (Scalar (Record Map Name (TypeBase dim u)
m2)) =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase dim u -> TypeBase dim u -> TypeBase dim u
comb Map Name (TypeBase dim u)
m1 Map Name (TypeBase dim u)
m2
comb (Scalar (Sum Map Name [TypeBase dim u]
m1)) (Scalar (Sum Map Name [TypeBase dim u]
m2)) =
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase dim u -> TypeBase dim u -> TypeBase dim u
comb) Map Name [TypeBase dim u]
m1 Map Name [TypeBase dim u]
m2
comb TypeBase dim u
t1 TypeBase dim u
_ = TypeBase dim u
t1
updatePat pat :: Pat ParamType
pat@(Wildcard (Info ParamType
tp) SrcLoc
loc) StaticVal
sv
| forall dim as. TypeBase dim as -> Bool
orderZero ParamType
tp = Pat ParamType
pat
| Bool
otherwise = forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) SrcLoc
loc
updatePat (PatAscription Pat ParamType
pat TypeExp Info VName
_ SrcLoc
_) StaticVal
sv =
Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
sv
updatePat p :: Pat ParamType
p@PatLit {} StaticVal
_ = Pat ParamType
p
updatePat pat :: Pat ParamType
pat@(PatConstr Name
c1 (Info ParamType
t) [Pat ParamType]
ps SrcLoc
loc) sv :: StaticVal
sv@(SumSV Name
_ [StaticVal]
svs [(Name, [ParamType])]
_)
| forall dim as. TypeBase dim as -> Bool
orderZero ParamType
t = Pat ParamType
pat
| Bool
otherwise = forall (f :: * -> *) vn t.
Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
PatConstr Name
c1 (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe TypeBase Exp Uniqueness
t') [Pat ParamType]
ps' SrcLoc
loc
where
t' :: TypeBase Exp Uniqueness
t' = StaticVal -> TypeBase Exp Uniqueness
resTypeFromSV StaticVal
sv
ps' :: [Pat ParamType]
ps' = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Pat ParamType -> StaticVal -> Pat ParamType
updatePat [Pat ParamType]
ps [StaticVal]
svs
updatePat (PatConstr Name
c1 Info ParamType
_ [Pat ParamType]
ps SrcLoc
loc) (Dynamic ParamType
t) =
forall (f :: * -> *) vn t.
Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
PatConstr Name
c1 (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe ParamType
t) [Pat ParamType]
ps SrcLoc
loc
updatePat Pat ParamType
pat (Dynamic ParamType
t) = Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat (ParamType -> StaticVal
svFromType ParamType
t)
updatePat Pat ParamType
pat (HoleSV StructType
t SrcLoc
_) = Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat (ParamType -> StaticVal
svFromType forall a b. (a -> b) -> a -> b
$ forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe StructType
t)
updatePat Pat ParamType
pat StaticVal
sv =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Tried to update pattern\n"
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Pat ParamType
pat
forall a. [a] -> [a] -> [a]
++ [Char]
"\nto reflect the static value\n"
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv
svFromType :: ParamType -> StaticVal
svFromType :: ParamType -> StaticVal
svFromType (Scalar (Record Map Name ParamType
fs)) = [(Name, StaticVal)] -> StaticVal
RecordSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map ParamType -> StaticVal
svFromType Map Name ParamType
fs
svFromType ParamType
t = ParamType -> StaticVal
Dynamic ParamType
t
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind (ValBind Maybe (Info EntryPoint)
entry VName
name Maybe (TypeExp Info VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc)
| Scalar Arrow {} <- forall dim as. RetTypeBase dim as -> TypeBase dim as
retType ResRetType
rettype = do
([Pat ParamType]
body_pats, Exp
body', ResRetType
rettype') <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a b. a -> b -> a
const forall a. Monoid a => a
mempty) ResRetType
rettype) Exp
body
ValBind -> DefM (ValBind, Env)
defuncValBind forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> [TypeParamBase vn]
-> [PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo vn]
-> SrcLoc
-> ValBindBase f vn
ValBind
Maybe (Info EntryPoint)
entry
VName
name
forall a. Maybe a
Nothing
(forall a. a -> Info a
Info ResRetType
rettype')
[TypeParamBase VName]
tparams
([Pat ParamType]
params forall a. Semigroup a => a -> a -> a
<> [Pat ParamType]
body_pats)
Exp
body'
forall a. Maybe a
Nothing
[AttrInfo VName]
attrs
SrcLoc
loc
defuncValBind valbind :: ValBind
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp Info VName)
retdecl (Info (RetType [VName]
ret_dims TypeBase Exp Uniqueness
rettype)) [TypeParamBase VName]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_) = do
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any forall vn. TypeParamBase vn -> Bool
isTypeParam [TypeParamBase VName]
tparams) forall a b. (a -> b) -> a -> b
$
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
forall a. Show a => a -> [Char]
show VName
name
forall a. [a] -> [a] -> [a]
++ [Char]
" has type parameters, "
forall a. [a] -> [a] -> [a]
++ [Char]
"but the defunctionaliser expects a monomorphic input program."
([VName]
tparams', [Pat ParamType]
params', Exp
body', StaticVal
sv, TypeBase Exp Uniqueness
sv_t) <-
[VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> DefM
([VName], [Pat ParamType], Exp, StaticVal, TypeBase Exp Uniqueness)
defuncLet (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams) [Pat ParamType]
params Exp
body forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims TypeBase Exp Uniqueness
rettype
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let bound_sizes :: Set VName
bound_sizes = forall a. Ord a => [a] -> Set a
S.fromList (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall t. Pat t -> [VName]
patNames [Pat ParamType]
params') forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams' forall a. Semigroup a => a -> a -> a
<> Set VName
globals
[Pat ParamType]
params'' <- forall (m :: * -> *).
MonadFreshNames m =>
[Pat ParamType] -> m [Pat ParamType]
instAnySizes [Pat ParamType]
params'
let rettype' :: TypeBase Exp Uniqueness
rettype' = forall as.
Monoid as =>
TypeBase Exp as -> TypeBase Exp as -> TypeBase Exp as
combineTypeShapes TypeBase Exp Uniqueness
rettype TypeBase Exp Uniqueness
sv_t
tparams'' :: [VName]
tparams'' = [VName]
tparams' forall a. [a] -> [a] -> [a]
++ Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
params''
ret_dims' :: [VName]
ret_dims' = forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Set VName
bound_sizes) forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars forall a b. (a -> b) -> a -> b
$ forall u. TypeBase Exp u -> FV
freeInType TypeBase Exp Uniqueness
rettype'
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( ValBind
valbind
{ valBindRetDecl :: Maybe (TypeExp Info VName)
valBindRetDecl = Maybe (TypeExp Info VName)
retdecl,
valBindRetType :: Info ResRetType
valBindRetType =
forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Pat ParamType]
params'
then forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims' forall a b. (a -> b) -> a -> b
$ TypeBase Exp Uniqueness
rettype' forall dim u1 u2. TypeBase dim u1 -> u2 -> TypeBase dim u2
`setUniqueness` Uniqueness
Nonunique
else forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims' TypeBase Exp Uniqueness
rettype',
valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) [VName]
tparams'',
valBindParams :: [Pat ParamType]
valBindParams = [Pat ParamType]
params'',
valBindBody :: Exp
valBindBody = Exp
body'
},
forall k a. k -> a -> Map k a
M.singleton VName
name forall a b. (a -> b) -> a -> b
$
Maybe ([VName], StructType) -> StaticVal -> Binding
Binding
(forall a. a -> Maybe a
Just (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName) (ValBind -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBind
valbind)))
StaticVal
sv
)
defuncVals :: [ValBind] -> DefM ()
defuncVals :: [ValBind] -> DefM ()
defuncVals [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
defuncVals (ValBind
valbind : [ValBind]
ds) = do
(ValBind
valbind', Env
env) <- ValBind -> DefM (ValBind, Env)
defuncValBind ValBind
valbind
ValBind -> DefM ()
addValBind ValBind
valbind'
let globals :: [VName]
globals = ValBind -> [VName]
valBindBound ValBind
valbind'
forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ forall a. [VName] -> DefM a -> DefM a
areGlobal [VName]
globals forall a b. (a -> b) -> a -> b
$ [ValBind] -> DefM ()
defuncVals [ValBind]
ds
{-# NOINLINE transformProg #-}
transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg [ValBind]
decs = forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
let ((), VNameSource
namesrc', [ValBind]
decs') = forall a. VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM VNameSource
namesrc forall a b. (a -> b) -> a -> b
$ [ValBind] -> DefM ()
defuncVals [ValBind]
decs
in ([ValBind]
decs', VNameSource
namesrc')