module Futhark.Internalise.Defunctionalise (transformProg) where
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, tails)
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)
import Language.Futhark
import Language.Futhark.Traversals
data StaticVal
= Dynamic PatType
|
LambdaSV Pat StructRetType Exp Env
| RecordSV [(Name, StaticVal)]
|
SumSV Name [StaticVal] [(Name, [PatType])]
|
DynamicFun (Exp, StaticVal) StaticVal
| IntrinsicSV
| HoleSV PatType 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 (Maybe ([VName], StructType)) 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)
bindingSV :: Binding -> StaticVal
bindingSV :: Binding -> StaticVal
bindingSV (Binding Maybe ([VName], StructType)
_ StaticVal
sv) = StaticVal
sv
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 Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim
where
onDim :: Size -> Size
onDim (NamedSize QualName VName
v) =
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') -> QualName VName -> Size
NamedSize QualName VName
v'
Just (SubstConst Int
d) -> Int -> Size
ConstSize Int
d
Maybe SizeSubst
Nothing -> QualName VName -> Size
NamedSize QualName VName
v
onDim Size
d = Size
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 PatBase Info VName
param (RetType [VName]
t_dims StructType
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 PatBase Info VName -> StructRetType -> Exp -> Env -> StaticVal
LambdaSV
(forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs PatBase Info VName
param)
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
t_dims (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs StructType
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 PatType
t ->
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
orig_substs PatType
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, [PatType])]
ts ->
Name -> [StaticVal] -> [(Name, [PatType])] -> 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 Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
orig_substs) [(Name, [PatType])]
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 PatType
t SrcLoc
loc ->
PatType -> SrcLoc -> StaticVal
HoleSV PatType
t SrcLoc
loc
where
tv :: Map VName SizeSubst -> ASTMapper m
tv Map VName SizeSubst
substs =
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper
{ mapOnPatType :: PatType -> m PatType
mapOnPatType = 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 Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs,
mapOnStructType :: StructType -> m 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 Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs,
mapOnExp :: Exp -> m 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 -> m 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 PatType
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 PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' Info PatType
t SrcLoc
loc
Just (SubstConst Int
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 Int
d))) SrcLoc
loc
Maybe SizeSubst
Nothing ->
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
v (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info PatType
t) SrcLoc
loc
onExp Map VName SizeSubst
substs (AppExp (Coerce Exp
e TypeExp VName
te SrcLoc
loc) (Info (AppRes PatType
t [VName]
ext))) =
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> AppExpBase f vn
Coerce (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e) TypeExp VName
te' SrcLoc
loc) (forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs PatType
t) [VName]
ext))
where
te' :: TypeExp VName
te' = Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te
onExp Map VName SizeSubst
substs (Lambda [PatBase Info VName]
params Exp
e Maybe (TypeExp VName)
ret (Info (Set Alias
als, RetType [VName]
t_dims StructType
t)) SrcLoc
loc) =
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda
(forall a b. (a -> b) -> [a] -> [b]
map (forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs) [PatBase Info VName]
params)
(Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e)
Maybe (TypeExp VName)
ret
(forall a. a -> Info a
Info (Set Alias
als, forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
t_dims (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs StructType
t)))
SrcLoc
loc
onExp Map VName SizeSubst
substs Exp
e = forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs Exp
e
onTypeExpDim :: Map VName SizeSubst -> SizeExp VName -> SizeExp VName
onTypeExpDim Map VName SizeSubst
substs d :: SizeExp VName
d@(SizeExpNamed QualName VName
v 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 vn. QualName vn -> SrcLoc -> SizeExp vn
SizeExpNamed QualName VName
v' SrcLoc
loc
Just (SubstConst Int
x) ->
forall vn. Int -> SrcLoc -> SizeExp vn
SizeExpConst Int
x SrcLoc
loc
Maybe SizeSubst
Nothing ->
SizeExp VName
d
onTypeExpDim Map VName SizeSubst
_ SizeExp VName
d = SizeExp VName
d
onTypeArgExp :: Map VName SizeSubst -> TypeArgExp VName -> TypeArgExp VName
onTypeArgExp Map VName SizeSubst
substs (TypeArgExpDim SizeExp VName
d SrcLoc
loc) =
forall vn. SizeExp vn -> SrcLoc -> TypeArgExp vn
TypeArgExpDim (Map VName SizeSubst -> SizeExp VName -> SizeExp VName
onTypeExpDim Map VName SizeSubst
substs SizeExp VName
d) SrcLoc
loc
onTypeArgExp Map VName SizeSubst
substs (TypeArgExpType TypeExp VName
te) =
forall vn. TypeExp vn -> TypeArgExp vn
TypeArgExpType (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te)
onTypeExp :: Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs (TEArray SizeExp VName
d TypeExp VName
te SrcLoc
loc) =
forall vn. SizeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArray (Map VName SizeSubst -> SizeExp VName -> SizeExp VName
onTypeExpDim Map VName SizeSubst
substs SizeExp VName
d) (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEUnique TypeExp VName
t SrcLoc
loc) =
forall vn. TypeExp vn -> SrcLoc -> TypeExp vn
TEUnique (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEApply TypeExp VName
t1 TypeArgExp VName
t2 SrcLoc
loc) =
forall vn. TypeExp vn -> TypeArgExp vn -> SrcLoc -> TypeExp vn
TEApply (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t1) (Map VName SizeSubst -> TypeArgExp VName -> TypeArgExp VName
onTypeArgExp Map VName SizeSubst
substs TypeArgExp VName
t2) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEArrow Maybe VName
p TypeExp VName
t1 TypeExp VName
t2 SrcLoc
loc) =
forall vn.
Maybe vn -> TypeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArrow Maybe VName
p (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t1) (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t2) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TETuple [TypeExp VName]
ts SrcLoc
loc) =
forall vn. [TypeExp vn] -> SrcLoc -> TypeExp vn
TETuple (forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [TypeExp VName]
ts) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TERecord [(Name, TypeExp VName)]
ts SrcLoc
loc) =
forall vn. [(Name, TypeExp vn)] -> SrcLoc -> TypeExp vn
TERecord (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
$ Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [(Name, TypeExp VName)]
ts) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TESum [(Name, [TypeExp VName])]
ts SrcLoc
loc) =
forall vn. [(Name, [TypeExp vn])] -> SrcLoc -> TypeExp vn
TESum (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
$ Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [(Name, [TypeExp VName])]
ts) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEDim [VName]
dims TypeExp VName
t SrcLoc
loc) =
forall vn. [vn] -> TypeExp vn -> SrcLoc -> TypeExp vn
TEDim [VName]
dims (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t) SrcLoc
loc
onTypeExp Map VName SizeSubst
_ (TEVar QualName VName
v SrcLoc
loc) =
forall vn. QualName vn -> SrcLoc -> TypeExp vn
TEVar QualName VName
v SrcLoc
loc
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 Size als -> TypeBase Size 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)
onAST :: ASTMappable x => M.Map VName SizeSubst -> x -> x
onAST :: forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (forall {m :: * -> *}. Monad m => Map VName SizeSubst -> ASTMapper m
tv Map VName SizeSubst
substs)
restrictEnvTo :: FV -> DefM Env
restrictEnvTo :: FV -> DefM Env
restrictEnvTo (FV Map VName StructType
m) = 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 forall a b. (a -> b) -> a -> b
$ VName
k forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals
Uniqueness
u <- forall shape as. TypeBase shape as -> Uniqueness
uniqueness 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
k Map VName StructType
m
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
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv
restrict' :: Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
Nonunique (Dynamic PatType
t) =
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ PatType
t forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
restrict' Uniqueness
_ (Dynamic PatType
t) =
PatType -> StaticVal
Dynamic PatType
t
restrict' Uniqueness
u (LambdaSV PatBase Info VName
pat StructRetType
t Exp
e Env
env) =
PatBase Info VName -> StructRetType -> Exp -> Env -> StaticVal
LambdaSV PatBase Info VName
pat StructRetType
t Exp
e forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Uniqueness -> Binding -> Binding
restrict'' Uniqueness
u) Env
env
restrict' Uniqueness
u (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 forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u) [(Name, StaticVal)]
fields
restrict' Uniqueness
u (SumSV Name
c [StaticVal]
svs [(Name, [PatType])]
fields) =
Name -> [StaticVal] -> [(Name, [PatType])] -> StaticVal
SumSV Name
c (forall a b. (a -> b) -> [a] -> [b]
map (Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u) [StaticVal]
svs) [(Name, [PatType])]
fields
restrict' Uniqueness
u (DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2) =
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp
e, Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv1) forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv2
restrict' Uniqueness
_ StaticVal
IntrinsicSV = StaticVal
IntrinsicSV
restrict' Uniqueness
_ (HoleSV PatType
t SrcLoc
loc) = PatType -> SrcLoc -> StaticVal
HoleSV PatType
t SrcLoc
loc
restrict'' :: Uniqueness -> Binding -> Binding
restrict'' Uniqueness
u (Binding Maybe ([VName], StructType)
t StaticVal
sv) = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u 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 :)
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
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
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 ()
_ Uniqueness
_ QualName VName
_ [TypeArg Size]
targs)) =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map TypeArg Size -> Set VName
f [TypeArg Size]
targs
where
f :: TypeArg Size -> Set VName
f (TypeArgDim (NamedSize QualName VName
d) 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 SrcLoc
_) = StructType -> Set VName
arraySizes StructType
t
arraySizes (Scalar Prim {}) = forall a. Monoid a => a
mempty
arraySizes (Array ()
_ Uniqueness
_ Shape Size
shape ScalarTypeBase Size ()
t) =
StructType -> Set VName
arraySizes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase Size ()
t) forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Size -> Set VName
dimName (forall dim. Shape dim -> [dim]
shapeDims Shape Size
shape)
where
dimName :: Size -> S.Set VName
dimName :: Size -> Set VName
dimName (NamedSize QualName VName
qn) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
dimName Size
_ = forall a. Monoid a => a
mempty
patternArraySizes :: Pat -> S.Set VName
patternArraySizes :: PatBase Info VName -> Set VName
patternArraySizes = StructType -> Set VName
arraySizes forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> StructType
patternStructType
data SizeSubst
= SubstNamed (QualName VName)
| SubstConst Int
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 Size a -> TypeBase Size a -> Map VName SizeSubst
dimMapping TypeBase Size a
t1 TypeBase Size 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 :: * -> *}.
(Foldable t, MonadState (Map VName SizeSubst) f) =>
t VName -> Size -> Size -> f Size
f TypeBase Size a
t1 TypeBase Size a
t2) forall a. Monoid a => a
mempty
where
f :: t VName -> Size -> Size -> f Size
f t VName
bound Size
d1 (NamedSize QualName VName
d2)
| 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 Size
d1
f t VName
_ (NamedSize QualName VName
d1) (NamedSize QualName VName
d2) = 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 VName
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
$ QualName VName -> Size
NamedSize QualName VName
d1
f t VName
_ (NamedSize QualName VName
d1) (ConstSize Int
d2) = 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 VName
d1) forall a b. (a -> b) -> a -> b
$ Int -> SizeSubst
SubstConst Int
d2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize QualName VName
d1
f t VName
_ Size
d Size
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d
dimMapping' ::
Monoid a =>
TypeBase Size a ->
TypeBase Size a ->
M.Map VName VName
dimMapping' :: forall a.
Monoid a =>
TypeBase Size a -> TypeBase Size a -> Map VName VName
dimMapping' TypeBase Size a
t1 TypeBase Size 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 Size a -> TypeBase Size a -> Map VName SizeSubst
dimMapping TypeBase Size a
t1 TypeBase Size 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, [PatType])]
_) =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StaticVal -> Set VName
sizesToRename [StaticVal]
svs
sizesToRename (LambdaSV PatBase Info VName
param StructRetType
_ Exp
_ Env
_) =
PatBase Info VName -> Set VName
freeInPat PatBase Info VName
param
forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall a. (a -> Bool) -> Set a -> Set a
S.filter forall {vn}. IdentBase Info vn -> Bool
couldBeSize forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents PatBase Info VName
param)
where
couldBeSize :: IdentBase Info vn -> Bool
couldBeSize IdentBase Info vn
ident =
forall a. Info a -> a
unInfo (forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType IdentBase Info vn
ident) forall a. Eq a => a -> a -> Bool
== forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim (IntType -> PrimType
Signed IntType
Int64))
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 a b. (a -> b) -> a -> b
$ 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 Size a -> TypeBase Size a -> Map VName SizeSubst
dimMapping (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size 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] ->
Exp ->
StructRetType ->
SrcLoc ->
DefM (Exp, StaticVal)
defuncFun :: [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
tparams [PatBase Info VName]
pats Exp
e0 StructRetType
ret SrcLoc
loc = do
let (PatBase Info VName
pat, StructRetType
ret', Exp
e0') = case [PatBase Info VName]
pats of
[] -> forall a. HasCallStack => [Char] -> a
error [Char]
"Received a lambda with no parameters."
[PatBase Info VName
pat'] -> (PatBase Info VName
pat', StructRetType
ret, Exp
e0)
(PatBase Info VName
pat' : [PatBase Info VName]
pats') ->
( PatBase Info VName
pat',
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType (forall a b. (a -> b) -> [a] -> [b]
map (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> PatType
patternType) [PatBase Info VName]
pats') StructRetType
ret,
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats' Exp
e0 forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
ret)) SrcLoc
loc
)
let used :: FV
used =
Exp -> FV
freeInExp (forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
e0 forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
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
. forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> PatType
typeFromSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV) Env
used_env
forall a. Semigroup a => a -> a -> a
<> PatBase Info VName -> Set VName
patternArraySizes PatBase Info VName
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,
PatBase Info VName -> StructRetType -> Exp -> Env -> StaticVal
LambdaSV PatBase Info VName
pat StructRetType
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' :: PatType
tp' = StaticVal -> PatType
typeFromSV StaticVal
sv
in ( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
Name
name
(forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
vn) (forall a. a -> Info a
Info PatType
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, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp e :: Exp
e@IntLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp e :: Exp
e@FloatLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp e :: Exp
e@StringLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
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 a b. [(a, b)] -> ([a], [b])
unzip 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 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 a b. [(a, b)] -> ([a], [b])
unzip 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 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 PatType
t) SrcLoc
loc') = do
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
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 PatType
tp = forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv
in forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
vn -> f PatType -> SrcLoc -> FieldBase f vn
RecordFieldImplicit VName
vn Info PatType
tp SrcLoc
loc', (VName -> Name
baseName VName
vn, StaticVal
sv))
defuncExp (ArrayLit [Exp]
es t :: Info PatType
t@(Info PatType
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 PatType -> SrcLoc -> ExpBase f vn
ArrayLit [Exp]
es' Info PatType
t SrcLoc
loc, PatType -> StaticVal
Dynamic PatType
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,
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ AppRes -> PatType
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 PatType
t) SrcLoc
loc) = do
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
case StaticVal
sv of
DynamicFun (Exp, StaticVal)
closure StaticVal
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp, StaticVal)
closure
StaticVal
IntrinsicSV -> do
([PatBase Info VName]
pats, Exp
body, StructRetType
tp) <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (Exp -> PatType
typeOf Exp
e) Exp
e
Exp -> DefM (Exp, StaticVal)
defuncExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
tp)) forall a. Monoid a => a
mempty
HoleSV PatType
_ SrcLoc
hole_loc ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. f PatType -> SrcLoc -> ExpBase f vn
Hole (forall a. a -> Info a
Info PatType
t) SrcLoc
hole_loc, StaticVal
sv)
StaticVal
_ ->
let tp :: PatType
tp = StaticVal -> PatType
typeFromSV StaticVal
sv
in forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info PatType
tp) SrcLoc
loc, StaticVal
sv)
defuncExp (Hole (Info PatType
t) SrcLoc
loc) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. f PatType -> SrcLoc -> ExpBase f vn
Hole (forall a. a -> Info a
Info PatType
t) SrcLoc
loc, PatType -> SrcLoc -> StaticVal
HoleSV PatType
t SrcLoc
loc)
defuncExp (Ascript Exp
e0 TypeExp VName
tydecl SrcLoc
loc)
| forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
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 vn -> SrcLoc -> ExpBase f vn
Ascript Exp
e0' TypeExp VName
tydecl SrcLoc
loc, StaticVal
sv)
| Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (AppExp (Coerce Exp
e0 TypeExp VName
tydecl SrcLoc
loc) Info AppRes
res)
| forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
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.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> AppExpBase f vn
Coerce Exp
e0' TypeExp VName
tydecl SrcLoc
loc) Info AppRes
res, StaticVal
sv)
| Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e1 Exp
e2 SrcLoc
loc) (Info (AppRes PatType
t [VName]
retext))) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
let env :: Env
env = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv1
pat' :: PatBase Info VName
pat' = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
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 Size a -> TypeBase Size a -> Map VName VName
dimMapping' (Exp -> PatType
typeOf Exp
e2) PatType
t
subst :: VName -> VName
subst VName
v = forall a. a -> Maybe a -> a
fromMaybe VName
v forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
mapping
mapper :: ASTMapper Identity
mapper = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnName :: VName -> Identity VName
mapOnName = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> VName
subst}
t' :: PatType
t' = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper Identity
mapper) forall a b. (a -> b) -> a -> b
$ Exp -> PatType
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
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes PatBase Info VName
pat' Exp
e1' Exp
e2' SrcLoc
loc) (forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
t' [VName]
retext)), StaticVal
sv2)
defuncExp (AppExp (LetFun VName
vn ([TypeParamBase VName], [PatBase Info VName],
Maybe (TypeExp VName), Info StructRetType, 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 e :: Exp
e@(AppExp (Apply f :: Exp
f@(Var QualName VName
f' Info PatType
_ SrcLoc
_) Exp
arg Info (Diet, Maybe VName)
d SrcLoc
loc) Info AppRes
res)
| VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
f') forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
TupLit [Exp]
es SrcLoc
tuploc <- Exp
arg = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncSoacExp [Exp]
es
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
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f (forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [Exp]
es' SrcLoc
tuploc) Info (Diet, Maybe VName)
d SrcLoc
loc) Info AppRes
res,
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e
)
defuncExp e :: Exp
e@(AppExp Apply {} Info AppRes
_) = Int -> Exp -> DefM (Exp, StaticVal)
defuncApply Int
0 Exp
e
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 [PatBase Info VName]
pats Exp
e0 Maybe (TypeExp VName)
_ (Info (Set Alias
_, StructRetType
ret)) SrcLoc
loc) =
[VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [PatBase Info VName]
pats Exp
e0 StructRetType
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 PatBase Info VName
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 = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv1
(LoopFormBase Info VName
form', Env
env2) <- case LoopFormBase Info VName
form of
For IdentBase Info VName
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 -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
v Exp
e2', forall {k}. IdentBase Info k -> Map k Binding
envFromIdent IdentBase Info VName
v)
ForIn PatBase Info VName
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 -> ExpBase f vn -> LoopFormBase f vn
ForIn PatBase Info VName
pat2 Exp
e2', PatBase Info VName -> Env
envFromPat PatBase Info VName
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
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop [VName]
sparams PatBase Info VName
pat Exp
e1' LoopFormBase Info VName
form' Exp
e3' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
where
envFromIdent :: IdentBase Info k -> Map k Binding
envFromIdent (Ident k
vn (Info PatType
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
$ PatType -> StaticVal
Dynamic PatType
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 PatType
tp@(Info PatType
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 PatType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) SrcLoc
loc, StaticVal
sv)
Maybe StaticVal
Nothing -> forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
Dynamic PatType
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' Info PatType
tp SrcLoc
loc, PatType -> StaticVal
Dynamic PatType
tp')
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
id1 IdentBase Info VName
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
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType IdentBase Info VName
id1
(Exp
body', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localEnv (forall k a. k -> a -> Map k a
M.singleton (forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
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
-> IdentBase f vn
-> SliceBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetWith IdentBase Info VName
id1 IdentBase Info VName
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,
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
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 PatType
_ 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 PatType -> 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 -> PatType
typeFromSV 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 :: PatType
t@(Scalar Record {})) StaticVal
sv2 fs' :: [Name]
fs'@(Name
_ : [Name]
_) =
StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (PatType -> StaticVal
svFromType PatType
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 :: PatType
sum_t@(Scalar (Sum Map Name [PatType]
all_fs))) SrcLoc
loc) = do
([Exp]
es', [StaticVal]
svs) <- forall a b. [(a, b)] -> ([a], [b])
unzip 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 Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
let sv :: StaticVal
sv =
Name -> [StaticVal] -> [(Name, [PatType])] -> 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 als. Monoid als => TypeBase Size als -> TypeBase Size als
defuncType) Map Name [PatType]
all_fs
sum_t' :: PatType
sum_t' = forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes PatType
sum_t (StaticVal -> PatType
typeFromSV StaticVal
sv)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
Constr Name
name [Exp]
es' (forall a. a -> Info a
Info PatType
sum_t') SrcLoc
loc, StaticVal
sv)
where
defuncType ::
Monoid als =>
TypeBase Size als ->
TypeBase Size als
defuncType :: forall als. Monoid als => TypeBase Size als -> TypeBase Size als
defuncType (Array als
as Uniqueness
u Shape Size
shape ScalarTypeBase Size ()
t) = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array als
as Uniqueness
u Shape Size
shape (forall als.
Monoid als =>
ScalarTypeBase Size als -> ScalarTypeBase Size als
defuncScalar ScalarTypeBase Size ()
t)
defuncType (Scalar ScalarTypeBase Size als
t) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall als.
Monoid als =>
ScalarTypeBase Size als -> ScalarTypeBase Size als
defuncScalar ScalarTypeBase Size als
t
defuncScalar ::
Monoid als =>
ScalarTypeBase Size als ->
ScalarTypeBase Size als
defuncScalar :: forall als.
Monoid als =>
ScalarTypeBase Size als -> ScalarTypeBase Size als
defuncScalar (Record Map Name (TypeBase Size als)
fs) = forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
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 Size als -> TypeBase Size als
defuncType Map Name (TypeBase Size als)
fs
defuncScalar Arrow {} = forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a. Monoid a => a
mempty
defuncScalar (Sum Map Name [TypeBase Size als]
fs) = forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
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 Size als -> TypeBase Size als
defuncType) Map Name [TypeBase Size als]
fs
defuncScalar (Prim PrimType
t) = forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
defuncScalar (TypeVar als
as Uniqueness
u QualName VName
tn [TypeArg Size]
targs) = forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar als
as Uniqueness
u QualName VName
tn [TypeArg Size]
targs
defuncExp (Constr Name
name [Exp]
_ (Info PatType
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 PatType
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
NonEmpty (CaseBase Info VName, StaticVal)
csPairs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StaticVal
-> CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal)
defuncCase StaticVal
sv) 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 (Case, StaticVal)
defuncCase :: StaticVal
-> CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal)
defuncCase StaticVal
sv (CasePat PatBase Info VName
p Exp
e SrcLoc
loc) = do
let p' :: PatBase Info VName
p' = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
p StaticVal
sv
env :: Env
env = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
p StaticVal
sv
(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 (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat PatBase Info VName
p' Exp
e' SrcLoc
loc, StaticVal
sv')
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 [PatBase Info VName]
params Exp
e0 Maybe (TypeExp VName)
decl Info (Set Alias, StructRetType)
tp SrcLoc
loc) = do
let env :: Env
env = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
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]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
params Exp
e0' Maybe (TypeExp VName)
decl Info (Set Alias, StructRetType)
tp SrcLoc
loc
defuncSoacExp Exp
e
| Scalar Arrow {} <- Exp -> PatType
typeOf Exp
e = do
([PatBase Info VName]
pats, Exp
body, StructRetType
tp) <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (Exp -> PatType
typeOf Exp
e) Exp
e
let env :: Env
env = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
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]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
body' forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
tp)) forall a. Monoid a => a
mempty
| Bool
otherwise = Exp -> DefM Exp
defuncExp' Exp
e
etaExpand :: PatType -> Exp -> DefM ([Pat], Exp, StructRetType)
etaExpand :: PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand PatType
e_t Exp
e = do
let ([(PName, StructType)]
ps, RetTypeBase Size (Set Alias)
ret) = forall {dim} {as}.
RetTypeBase dim as
-> ([(PName, TypeBase dim ())], RetTypeBase dim as)
getType forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
e_t
([VName]
_, ([PatBase Info VName]
pats, [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 :: * -> *) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM forall {m :: * -> *} {as}.
MonadFreshNames m =>
[VName]
-> (PName, TypeBase Size as)
-> m ([VName], (PatBase Info VName, Exp))
f [] [(PName, StructType)]
ps
let e' :: Exp
e' =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
( \Exp
e1 (Exp
e2, StructType
t2, [StructType]
argtypes) ->
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
e1 Exp
e2 (forall a. a -> Info a
Info (forall shape as. TypeBase shape as -> Diet
diet StructType
t2, forall a. Maybe a
Nothing)) forall a. Monoid a => a
mempty)
(forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [StructType]
argtypes RetTypeBase Size (Set Alias)
ret) []))
)
Exp
e
forall a b. (a -> b) -> a -> b
$ forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Exp]
vars (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(PName, StructType)]
ps) (forall a. Int -> [a] -> [a]
drop Int
1 forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [[a]]
tails forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(PName, StructType)]
ps)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PatBase Info VName]
pats, Exp
e', forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a b. a -> b -> a
const ()) RetTypeBase Size (Set Alias)
ret)
where
getType :: RetTypeBase dim as
-> ([(PName, TypeBase dim ())], RetTypeBase dim as)
getType (RetType [VName]
_ (Scalar (Arrow as
_ PName
p TypeBase dim ()
t1 RetTypeBase dim as
t2))) =
let ([(PName, TypeBase dim ())]
ps, RetTypeBase dim as
r) = RetTypeBase dim as
-> ([(PName, TypeBase dim ())], RetTypeBase dim as)
getType RetTypeBase dim as
t2 in ((PName
p, TypeBase dim ()
t1) forall a. a -> [a] -> [a]
: [(PName, TypeBase dim ())]
ps, RetTypeBase dim as
r)
getType RetTypeBase dim as
t = ([], RetTypeBase dim as
t)
f :: [VName]
-> (PName, TypeBase Size as)
-> m ([VName], (PatBase Info VName, Exp))
f [VName]
prev (PName
p, TypeBase Size as
t) = do
let t' :: PatType
t' = forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase Size as
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]
"x"
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( VName
x forall a. a -> [a] -> [a]
: [VName]
prev,
( forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
x (forall a. a -> Info a
Info PatType
t') forall a. Monoid a => a
mempty,
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
x) (forall a. a -> Info a
Info PatType
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] ->
Exp ->
StructRetType ->
DefM ([VName], [Pat], Exp, StaticVal)
defuncLet :: [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> DefM ([VName], [PatBase Info VName], Exp, StaticVal)
defuncLet [VName]
dims ps :: [PatBase Info VName]
ps@(PatBase Info VName
pat : [PatBase Info VName]
pats) Exp
body (RetType [VName]
ret_dims StructType
rettype)
| forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info VName
pat = do
let bound_by_pat :: VName -> Bool
bound_by_pat = (forall a. Ord a => a -> Set a -> Bool
`S.member` PatBase Info VName -> Set VName
freeInPat PatBase Info VName
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 = PatBase Info VName -> Env
envFromPat PatBase Info VName
pat forall a. Semigroup a => a -> a -> a
<> [VName] -> Env
envFromDimNames [VName]
pat_dims
([VName]
rest_dims', [PatBase Info VName]
pats', Exp
body', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> DefM ([VName], [PatBase Info VName], Exp, StaticVal)
defuncLet [VName]
rest_dims [PatBase Info VName]
pats Exp
body forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
rettype
(Exp, StaticVal)
closure <- [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [PatBase Info VName]
ps Exp
body (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
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',
PatBase Info VName
pat forall a. a -> [a] -> [a]
: [PatBase Info VName]
pats',
Exp
body',
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp, StaticVal)
closure StaticVal
sv
)
| Bool
otherwise = do
(Exp
e, StaticVal
sv) <- [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [PatBase Info VName]
ps Exp
body (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
rettype) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], Exp
e, StaticVal
sv)
defuncLet [VName]
_ [] Exp
body (RetType [VName]
_ StructType
rettype) = do
(Exp
body', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], Exp
body', forall {as}. StaticVal -> TypeBase Size as -> StaticVal
imposeType StaticVal
sv StructType
rettype)
where
imposeType :: StaticVal -> TypeBase Size as -> StaticVal
imposeType Dynamic {} TypeBase Size as
t =
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase Size as
t
imposeType (RecordSV [(Name, StaticVal)]
fs1) (Scalar (Record Map Name (TypeBase Size as)
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 -> TypeBase Size as -> StaticVal
imposeType (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, StaticVal)]
fs1) Map Name (TypeBase Size as)
fs2
imposeType StaticVal
sv TypeBase Size as
_ = StaticVal
sv
sizesForAll :: MonadFreshNames m => S.Set VName -> [Pat] -> m ([VName], [Pat])
sizesForAll :: forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
params = do
([PatBase Info VName]
params', Set VName
sizes) <- forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT (Set VName) m)
tv) [PatBase Info VName]
params) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Set a -> [a]
S.toList Set VName
sizes, [PatBase Info VName]
params')
where
bound :: Set VName
bound = Set VName
bound_sizes forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase Info VName]
params
tv :: ASTMapper (StateT (Set VName) m)
tv = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatType :: PatType -> StateT (Set VName) m PatType
mapOnPatType = 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 {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadState (Set VName) (t m), MonadTrans t, MonadFreshNames m) =>
Size -> t m Size
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure}
onDim :: Size -> t m Size
onDim (AnySize (Just VName
v)) = do
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> Set a -> Set a
S.insert VName
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
v
onDim (AnySize Maybe VName
Nothing) = do
VName
v <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"size"
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> Set a -> Set a
S.insert VName
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
v
onDim (NamedSize QualName VName
d) = 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 a b. (a -> b) -> a -> b
$
forall a. Ord a => a -> Set a -> Set a
S.insert forall a b. (a -> b) -> a -> b
$
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
$ QualName VName -> Size
NamedSize QualName VName
d
onDim Size
d = forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d
unRetType :: StructRetType -> StructType
unRetType :: StructRetType -> StructType
unRetType (RetType [] StructType
t) = StructType
t
unRetType (RetType [VName]
ext StructType
t) = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim StructType
t
where
onDim :: Size -> Size
onDim (NamedSize QualName VName
d) | forall vn. QualName vn -> vn
qualLeaf QualName VName
d forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
ext = Maybe VName -> Size
AnySize forall a. Maybe a
Nothing
onDim Size
d = Size
d
defuncApply :: Int -> Exp -> DefM (Exp, StaticVal)
defuncApply :: Int -> Exp -> DefM (Exp, StaticVal)
defuncApply Int
depth e :: Exp
e@(AppExp (Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
d SrcLoc
loc) t :: Info AppRes
t@(Info (AppRes PatType
ret [VName]
ext))) = do
let ([StructType]
argtypes, StructType
_) = forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType PatType
ret
(Exp
e1', StaticVal
sv1) <- Int -> Exp -> DefM (Exp, StaticVal)
defuncApply (Int
depth forall a. Num a => a -> a -> a
+ Int
1) Exp
e1
(Exp
e2', StaticVal
sv2) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
let e' :: Exp
e' = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
e1' Exp
e2' Info (Diet, Maybe VName)
d SrcLoc
loc) Info AppRes
t
case StaticVal
sv1 of
LambdaSV PatBase Info VName
pat StructRetType
e0_t Exp
e0 Env
closure_env -> do
let env' :: Env
env' = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv2
dims :: [VName]
dims = forall a. Monoid a => a
mempty
(Exp
e0', 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
e0
let closure_pat :: PatBase Info VName
closure_pat = [VName] -> Env -> PatBase Info VName
buildEnvPat [VName]
dims Env
closure_env
pat' :: PatBase Info VName
pat' = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv2
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let params :: [PatBase Info VName]
params = [PatBase Info VName
closure_pat, PatBase Info VName
pat']
params_for_rettype :: [PatBase Info VName]
params_for_rettype = [PatBase Info VName]
params forall a. [a] -> [a] -> [a]
++ StaticVal -> [PatBase Info VName]
svParams StaticVal
sv1 forall a. [a] -> [a] -> [a]
++ StaticVal -> [PatBase Info VName]
svParams StaticVal
sv2
svParams :: StaticVal -> [PatBase Info VName]
svParams (LambdaSV PatBase Info VName
sv_pat StructRetType
_ Exp
_ Env
_) = [PatBase Info VName
sv_pat]
svParams StaticVal
_ = []
lifted_rettype :: PatType
lifted_rettype = Env -> [PatBase Info VName] -> StructType -> PatType -> PatType
buildRetType Env
closure_env [PatBase Info VName]
params_for_rettype (StructRetType -> StructType
unRetType StructRetType
e0_t) forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e0'
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 b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
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 PatBase Info VName -> Set VName
patternArraySizes [PatBase Info VName]
params
liftedName :: t -> ExpBase f VName -> [Char]
liftedName t
i (Var QualName VName
f f PatType
_ SrcLoc
_) =
[Char]
"defunc_" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show t
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 t
i (AppExp (Apply ExpBase f VName
f ExpBase f VName
_ f (Diet, Maybe VName)
_ SrcLoc
_) f AppRes
_) =
t -> ExpBase f VName -> [Char]
liftedName (t
i forall a. Num a => a -> a -> a
+ t
1) ExpBase f VName
f
liftedName t
_ ExpBase f VName
_ = [Char]
"defunc"
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
([VName]
missing_dims, [PatBase Info VName]
params') <- forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
params
VName
fname <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString forall a b. (a -> b) -> a -> b
$ forall {t} {f :: * -> *}.
(Show t, Num t) =>
t -> ExpBase f VName -> [Char]
liftedName (Int
0 :: Int) Exp
e1
VName
-> StructRetType
-> [VName]
-> [PatBase Info VName]
-> Exp
-> DefM ()
liftValDec
VName
fname
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
lifted_rettype)
([VName]
dims forall a. [a] -> [a] -> [a]
++ [VName]
more_dims forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims)
[PatBase Info VName]
params'
Exp
e0'
let t1 :: StructType
t1 = forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e1'
t2 :: StructType
t2 = forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e2'
fname' :: QualName VName
fname' = forall v. v -> QualName v
qualName VName
fname
fname'' :: Exp
fname'' =
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
QualName VName
fname'
( forall a. a -> Info a
Info
( forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed StructType
t1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed StructType
t2 forall a b. (a -> b) -> a -> b
$
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
lifted_rettype
)
)
SrcLoc
loc
callret :: AppRes
callret = PatType -> [VName] -> AppRes
AppRes (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes PatType
ret PatType
lifted_rettype) [VName]
ext
innercallret :: AppRes
innercallret =
PatType -> [VName] -> AppRes
AppRes
(forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed StructType
t2 forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
lifted_rettype)
[]
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
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
fname'' Exp
e1' (forall a. a -> Info a
Info (Diet
Observe, forall a. Maybe a
Nothing)) SrcLoc
loc)
(forall a. a -> Info a
Info AppRes
innercallret)
)
Exp
e2'
Info (Diet, Maybe VName)
d
SrcLoc
loc
)
(forall a. a -> Info a
Info AppRes
callret),
StaticVal
sv
)
DynamicFun (Exp, StaticVal)
_ StaticVal
sv -> do
let ([PatType]
argtypes', PatType
rettype) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv [StructType]
argtypes
restype :: PatType
restype = forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [PatType]
argtypes' (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
rettype) forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` forall as shape. Monoid as => TypeBase shape as -> as
aliases PatType
ret
callret :: AppRes
callret = PatType -> [VName] -> AppRes
AppRes (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes PatType
ret PatType
restype) [VName]
ext
apply_e :: Exp
apply_e = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
e1' Exp
e2' Info (Diet, Maybe VName)
d SrcLoc
loc) (forall a. a -> Info a
Info AppRes
callret)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
apply_e, StaticVal
sv)
StaticVal
IntrinsicSV -> forall {t :: * -> *} {a}.
Foldable t =>
t a -> Exp -> StaticVal -> DefM (Exp, StaticVal)
intrinsicOrHole [StructType]
argtypes Exp
e' StaticVal
sv1
HoleSV {} -> forall {t :: * -> *} {a}.
Foldable t =>
t a -> Exp -> StaticVal -> DefM (Exp, StaticVal)
intrinsicOrHole [StructType]
argtypes Exp
e' StaticVal
sv1
StaticVal
_ ->
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Application of an expression\n"
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Exp
e1
forall a. [a] -> [a] -> [a]
++ [Char]
"\nthat is neither a static lambda "
forall a. [a] -> [a] -> [a]
++ [Char]
"nor a dynamic function, but has static value:\n"
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv1
where
intrinsicOrHole :: t a -> Exp -> StaticVal -> DefM (Exp, StaticVal)
intrinsicOrHole t a
argtypes Exp
e' StaticVal
sv
| Int
depth forall a. Eq a => a -> a -> Bool
== Int
0 =
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null t a
argtypes
then forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e', PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
else do
([PatBase Info VName]
pats, Exp
body, StructRetType
tp) <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (Exp -> PatType
typeOf Exp
e') Exp
e'
Exp -> DefM (Exp, StaticVal)
defuncExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
tp)) forall a. Monoid a => a
mempty
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e', StaticVal
sv)
defuncApply Int
depth e :: Exp
e@(Var QualName VName
qn (Info PatType
t) SrcLoc
loc) = do
let ([StructType]
argtypes, StructType
_) = forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType PatType
t
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
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
depth -> do
let ([PatType]
argtypes', PatType
rettype) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv [StructType]
argtypes
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [PatType]
argtypes' forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
rettype)) SrcLoc
loc, StaticVal
sv)
| 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 ([PatBase Info VName]
pats, Exp
e0, StaticVal
sv') = [Char]
-> StaticVal -> Int -> ([PatBase Info VName], Exp, StaticVal)
liftDynFun (forall a. Pretty a => a -> [Char]
prettyString QualName VName
qn) StaticVal
sv Int
depth
([PatType]
argtypes', PatType
rettype) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv' [StructType]
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
([VName]
missing_dims, [PatBase Info VName]
pats') <- forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
pats
VName
-> StructRetType
-> [VName]
-> [PatBase Info VName]
-> Exp
-> DefM ()
liftValDec VName
fname (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
rettype) ([VName]
dims' forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims) [PatBase Info VName]
pats' Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
(forall v. v -> QualName v
qualName VName
fname)
(forall a. a -> Info a
Info (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [PatType]
argtypes' forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct PatType
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 PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info (StaticVal -> PatType
typeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncApply Int
depth (Parens Exp
e SrcLoc
_) = Int -> Exp -> DefM (Exp, StaticVal)
defuncApply Int
depth Exp
e
defuncApply Int
_ Exp
expr = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
expr
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], Exp, StaticVal)
liftDynFun :: [Char]
-> StaticVal -> Int -> ([PatBase Info VName], 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 PatBase Info VName
pat StructRetType
_ Exp
_ Env
_) StaticVal
sv) Int
d
| Int
d forall a. Ord a => a -> a -> Bool
> Int
0 =
let ([PatBase Info VName]
pats, Exp
e', StaticVal
sv') = [Char]
-> StaticVal -> Int -> ([PatBase Info VName], Exp, StaticVal)
liftDynFun [Char]
s StaticVal
sv (Int
d forall a. Num a => a -> a -> a
- Int
1)
in (PatBase Info VName
pat forall a. a -> [a] -> [a]
: [PatBase Info VName]
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 -> Env
envFromPat :: PatBase Info VName -> Env
envFromPat PatBase Info VName
pat = case PatBase Info VName
pat of
TuplePat [PatBase Info VName]
ps SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
ps
RecordPat [(Name, PatBase Info VName)]
fs SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatBase Info VName -> Env
envFromPat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, PatBase Info VName)]
fs
PatParens PatBase Info VName
p SrcLoc
_ -> PatBase Info VName -> Env
envFromPat PatBase Info VName
p
PatAttr AttrInfo VName
_ PatBase Info VName
p SrcLoc
_ -> PatBase Info VName -> Env
envFromPat PatBase Info VName
p
Id VName
vn (Info PatType
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
$ PatType -> StaticVal
Dynamic PatType
t
Wildcard Info PatType
_ SrcLoc
_ -> forall a. Monoid a => a
mempty
PatAscription PatBase Info VName
p TypeExp VName
_ SrcLoc
_ -> PatBase Info VName -> Env
envFromPat PatBase Info VName
p
PatLit {} -> forall a. Monoid a => a
mempty
PatConstr Name
_ Info PatType
_ [PatBase Info VName]
ps SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
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
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
liftValDec :: VName -> StructRetType -> [VName] -> [Pat] -> Exp -> DefM ()
liftValDec :: VName
-> StructRetType
-> [VName]
-> [PatBase Info VName]
-> Exp
-> DefM ()
liftValDec VName
fname (RetType [VName]
ret_dims StructType
ret) [VName]
dims [PatBase Info VName]
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 [VName]
dims forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
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 :: StructRetType
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 as. TypeBase Size as -> Set VName
freeInType StructType
ret)) forall a. [a] -> [a] -> [a]
++ [VName]
ret_dims) StructType
ret
dec :: ValBind
dec =
ValBind
{ valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = forall a. Maybe a
Nothing,
valBindName :: VName
valBindName = VName
fname,
valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = forall a. Maybe a
Nothing,
valBindRetType :: Info StructRetType
valBindRetType = forall a. a -> Info a
Info StructRetType
rettype_st,
valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
dims',
valBindParams :: [PatBase Info VName]
valBindParams = [PatBase Info VName]
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
}
buildEnvPat :: [VName] -> Env -> Pat
buildEnvPat :: [VName] -> Env -> PatBase Info VName
buildEnvPat [VName]
sizes Env
env = forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat (forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (Name, PatBase Info VName)
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, PatBase Info VName)
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. f PatType -> SrcLoc -> PatBase f vn
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) forall a. Monoid a => a
mempty
else forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
vn (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) forall a. Monoid a => a
mempty
)
buildRetType :: Env -> [Pat] -> StructType -> PatType -> PatType
buildRetType :: Env -> [PatBase Info VName] -> StructType -> PatType -> PatType
buildRetType Env
env [PatBase Info VName]
pats = forall {t :: * -> *} {shape} {as}.
(Foldable t, Monoid (t Alias)) =>
TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb
where
bound :: Set VName
bound =
forall a. Ord a => [a] -> Set a
S.fromList (forall k a. Map k a -> [k]
M.keys Env
env) forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
pats)
boundAsUnique :: VName -> Bool
boundAsUnique VName
v =
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall dim as. TypeBase dim as -> Bool
unique forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Info a -> a
unInfo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== VName
v) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) vn. IdentBase f vn -> vn
identName) forall a b. (a -> b) -> a -> b
$
forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
pats
problematic :: VName -> Bool
problematic VName
v = (VName
v forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) Bool -> Bool -> Bool
&& Bool -> Bool
not (VName -> Bool
boundAsUnique VName
v)
comb :: TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb (Scalar (Record Map Name (TypeBase shape as)
fs_annot)) (Scalar (Record Map Name (TypeBase shape (t Alias))
fs_got)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
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 shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb Map Name (TypeBase shape as)
fs_annot Map Name (TypeBase shape (t Alias))
fs_got
comb (Scalar (Sum Map Name [TypeBase shape as]
cs_annot)) (Scalar (Sum Map Name [TypeBase shape (t Alias)]
cs_got)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
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 shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb) Map Name [TypeBase shape as]
cs_annot Map Name [TypeBase shape (t Alias)]
cs_got
comb (Scalar Arrow {}) TypeBase shape (t Alias)
t =
forall {t :: * -> *} {dim}.
(Foldable t, Monoid (t Alias)) =>
TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend TypeBase shape (t Alias)
t
comb TypeBase shape as
got TypeBase shape (t Alias)
et =
forall {t :: * -> *} {dim}.
(Foldable t, Monoid (t Alias)) =>
TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase shape as
got forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase shape (t Alias)
et
descend :: TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend t :: TypeBase dim (t Alias)
t@Array {}
| forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Bool
problematic forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim (t Alias)
t) = TypeBase dim (t Alias)
t forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
descend (Scalar (Record Map Name (TypeBase dim (t Alias))
t)) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend Map Name (TypeBase dim (t Alias))
t
descend TypeBase dim (t Alias)
t = TypeBase dim (t Alias)
t
typeFromSV :: StaticVal -> PatType
typeFromSV :: StaticVal -> PatType
typeFromSV (Dynamic PatType
tp) =
PatType
tp
typeFromSV (LambdaSV PatBase Info VName
_ StructRetType
_ Exp
_ Env
env) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
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 -> PatType
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, PatType)]
ts = forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> PatType
typeFromSV) [(Name, StaticVal)]
ls
in forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatType)]
ts
typeFromSV (DynamicFun (Exp
_, StaticVal
sv) StaticVal
_) =
StaticVal -> PatType
typeFromSV StaticVal
sv
typeFromSV (SumSV Name
name [StaticVal]
svs [(Name, [PatType])]
fields) =
let svs' :: [PatType]
svs' = forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> PatType
typeFromSV [StaticVal]
svs
in forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
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 [PatType]
svs' forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [PatType])]
fields
typeFromSV (HoleSV PatType
t SrcLoc
_) =
PatType
t
typeFromSV StaticVal
IntrinsicSV =
forall a. HasCallStack => [Char] -> a
error [Char]
"Tried to get the type from the static value of an intrinsic."
dynamicFunType :: StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType :: StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) (StructType
p : [StructType]
ps) =
let ([PatType]
ps', PatType
ret) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv [StructType]
ps in (forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
p forall a. a -> [a] -> [a]
: [PatType]
ps', PatType
ret)
dynamicFunType StaticVal
sv [StructType]
_ = ([], StaticVal -> PatType
typeFromSV StaticVal
sv)
matchPatSV :: PatBase Info VName -> StaticVal -> Env
matchPatSV :: PatBase Info VName -> StaticVal -> Env
matchPatSV (TuplePat [PatBase Info VName]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls) =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\PatBase Info VName
p (Name
_, StaticVal
sv) -> PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
p StaticVal
sv) [PatBase Info VName]
ps [(Name, StaticVal)]
ls
matchPatSV (RecordPat [(Name, PatBase Info VName)]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls)
| [(Name, PatBase Info VName)]
ps' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, PatBase Info VName)]
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, PatBase Info VName)]
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 a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
_, PatBase Info VName
p) (Name
_, StaticVal
sv) -> PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
p StaticVal
sv) [(Name, PatBase Info VName)]
ps' [(Name, StaticVal)]
ls'
matchPatSV (PatParens PatBase Info VName
pat SrcLoc
_) StaticVal
sv = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv
matchPatSV (PatAttr AttrInfo VName
_ PatBase Info VName
pat SrcLoc
_) StaticVal
sv = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv
matchPatSV (Id VName
vn (Info PatType
t) SrcLoc
_) StaticVal
sv =
if StaticVal -> Bool
orderZeroSV StaticVal
sv
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
$ PatType -> StaticVal
Dynamic PatType
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
$ forall as. TypeBase Size as -> Set VName
freeInType PatType
t
i64 :: Binding
i64 = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
matchPatSV (Wildcard Info PatType
_ SrcLoc
_) StaticVal
_ = forall a. Monoid a => a
mempty
matchPatSV (PatAscription PatBase Info VName
pat TypeExp VName
_ SrcLoc
_) StaticVal
sv = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv
matchPatSV PatLit {} StaticVal
_ = forall a. Monoid a => a
mempty
matchPatSV (PatConstr Name
c1 Info PatType
_ [PatBase Info VName]
ps SrcLoc
_) (SumSV Name
c2 [StaticVal]
ls [(Name, [PatType])]
fs)
| Name
c1 forall a. Eq a => a -> a -> Bool
== Name
c2 =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> Env
matchPatSV [PatBase Info VName]
ps [StaticVal]
ls
| Just [PatType]
ts <- forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
c1 [(Name, [PatType])]
fs =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> Env
matchPatSV [PatBase Info VName]
ps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map PatType -> StaticVal
svFromType [PatType]
ts
| 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 PatType
_ [PatBase Info VName]
ps SrcLoc
_) (Dynamic (Scalar (Sum Map Name [PatType]
fs)))
| Just [PatType]
ts <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c1 Map Name [PatType]
fs =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> Env
matchPatSV [PatBase Info VName]
ps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map PatType -> StaticVal
svFromType [PatType]
ts
| 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 PatBase Info VName
pat (Dynamic PatType
t) = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
svFromType PatType
t
matchPatSV PatBase Info VName
pat (HoleSV PatType
t SrcLoc
_) = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
svFromType PatType
t
matchPatSV PatBase Info VName
pat StaticVal
sv =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Tried to match pattern "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatBase Info VName
pat
forall a. [a] -> [a] -> [a]
++ [Char]
" with static value "
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv
forall a. [a] -> [a] -> [a]
++ [Char]
"."
orderZeroSV :: StaticVal -> Bool
orderZeroSV :: StaticVal -> Bool
orderZeroSV Dynamic {} = Bool
True
orderZeroSV (RecordSV [(Name, StaticVal)]
fields) = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (StaticVal -> Bool
orderZeroSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, StaticVal)]
fields
orderZeroSV StaticVal
_ = Bool
False
updatePat :: Pat -> StaticVal -> Pat
updatePat :: PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat (TuplePat [PatBase Info VName]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs) =
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat [PatBase Info VName]
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, PatBase Info VName)]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs)
| [(Name, PatBase Info VName)]
ps' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, PatBase Info VName)]
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.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat
(forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
n, PatBase Info VName
p) (Name
_, StaticVal
sv) -> (Name
n, PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
p StaticVal
sv)) [(Name, PatBase Info VName)]
ps' [(Name, StaticVal)]
svs')
SrcLoc
loc
updatePat (PatParens PatBase Info VName
pat SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens (PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv) SrcLoc
loc
updatePat (PatAttr AttrInfo VName
attr PatBase Info VName
pat SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr (PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv) SrcLoc
loc
updatePat (Id VName
vn (Info PatType
tp) SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
vn (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall {dim} {as}.
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb PatType
tp (StaticVal -> PatType
typeFromSV StaticVal
sv forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique)) SrcLoc
loc
where
comb :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb (Scalar Arrow {}) TypeBase dim as
t2 = TypeBase dim as
t2
comb (Scalar (Record Map Name (TypeBase dim as)
m1)) (Scalar (Record Map Name (TypeBase dim as)
m2)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
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 as -> TypeBase dim as -> TypeBase dim as
comb Map Name (TypeBase dim as)
m1 Map Name (TypeBase dim as)
m2
comb (Scalar (Sum Map Name [TypeBase dim as]
m1)) (Scalar (Sum Map Name [TypeBase dim as]
m2)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
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 as -> TypeBase dim as -> TypeBase dim as
comb) Map Name [TypeBase dim as]
m1 Map Name [TypeBase dim as]
m2
comb TypeBase dim as
t1 TypeBase dim as
_ = TypeBase dim as
t1
updatePat pat :: PatBase Info VName
pat@(Wildcard (Info PatType
tp) SrcLoc
loc) StaticVal
sv
| forall dim as. TypeBase dim as -> Bool
orderZero PatType
tp = PatBase Info VName
pat
| Bool
otherwise = forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) SrcLoc
loc
updatePat (PatAscription PatBase Info VName
pat TypeExp VName
_ SrcLoc
_) StaticVal
sv =
PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv
updatePat p :: PatBase Info VName
p@PatLit {} StaticVal
_ = PatBase Info VName
p
updatePat pat :: PatBase Info VName
pat@(PatConstr Name
c1 (Info PatType
t) [PatBase Info VName]
ps SrcLoc
loc) sv :: StaticVal
sv@(SumSV Name
_ [StaticVal]
svs [(Name, [PatType])]
_)
| forall dim as. TypeBase dim as -> Bool
orderZero PatType
t = PatBase Info VName
pat
| Bool
otherwise = forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
c1 (forall a. a -> Info a
Info PatType
t') [PatBase Info VName]
ps' SrcLoc
loc
where
t' :: PatType
t' = StaticVal -> PatType
typeFromSV StaticVal
sv forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
ps' :: [PatBase Info VName]
ps' = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat [PatBase Info VName]
ps [StaticVal]
svs
updatePat (PatConstr Name
c1 Info PatType
_ [PatBase Info VName]
ps SrcLoc
loc) (Dynamic PatType
t) =
forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
c1 (forall a. a -> Info a
Info PatType
t) [PatBase Info VName]
ps SrcLoc
loc
updatePat PatBase Info VName
pat (Dynamic PatType
t) = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat (PatType -> StaticVal
svFromType PatType
t)
updatePat PatBase Info VName
pat (HoleSV PatType
t SrcLoc
_) = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat (PatType -> StaticVal
svFromType PatType
t)
updatePat PatBase Info VName
pat StaticVal
sv =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Tried to update pattern "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatBase Info VName
pat
forall a. [a] -> [a] -> [a]
++ [Char]
"to reflect the static value "
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv
svFromType :: PatType -> StaticVal
svFromType :: PatType -> StaticVal
svFromType (Scalar (Record Map Name PatType
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 PatType -> StaticVal
svFromType Map Name PatType
fs
svFromType PatType
t = PatType -> StaticVal
Dynamic PatType
t
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind (ValBind Maybe (Info EntryPoint)
entry VName
name Maybe (TypeExp VName)
_ (Info (RetType [VName]
_ StructType
rettype)) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc)
| Scalar Arrow {} <- StructType
rettype = do
([PatBase Info VName]
body_pats, Exp
body', StructRetType
rettype') <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
rettype) Exp
body
ValBind -> DefM (ValBind, Env)
defuncValBind forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f StructRetType
-> [TypeParamBase vn]
-> [PatBase f vn]
-> 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 StructRetType
rettype')
[TypeParamBase VName]
tparams
([PatBase Info VName]
params forall a. Semigroup a => a -> a -> a
<> [PatBase Info VName]
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 VName)
retdecl (Info (RetType [VName]
ret_dims StructType
rettype)) [TypeParamBase VName]
tparams [PatBase Info VName]
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', [PatBase Info VName]
params', Exp
body', StaticVal
sv) <-
[VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> DefM ([VName], [PatBase Info VName], Exp, StaticVal)
defuncLet (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams) [PatBase Info VName]
params Exp
body forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
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 (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase Info VName]
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
rettype' :: StructType
rettype' =
forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes StructType
rettype forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Set VName -> Size -> Size
anyDimIfNotBound Set VName
bound_sizes) forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
body'
ret_dims' :: [VName]
ret_dims' = forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Ord a => a -> Set a -> Bool
`S.member` forall as. TypeBase Size as -> Set VName
freeInType StructType
rettype') [VName]
ret_dims
([VName]
missing_dims, [PatBase Info VName]
params'') <- forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
params'
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( ValBind
valbind
{ valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
retdecl,
valBindRetType :: Info StructRetType
valBindRetType =
forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatBase Info VName]
params'
then forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims' forall a b. (a -> b) -> a -> b
$ StructType
rettype' forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
else forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims' StructType
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) forall a b. (a -> b) -> a -> b
$ [VName]
tparams' forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims,
valBindParams :: [PatBase Info VName]
valBindParams = [PatBase Info VName]
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
)
where
anyDimIfNotBound :: Set VName -> Size -> Size
anyDimIfNotBound Set VName
bound_sizes (NamedSize QualName VName
v)
| forall vn. QualName vn -> vn
qualLeaf QualName VName
v forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
bound_sizes = Maybe VName -> Size
AnySize forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
v
anyDimIfNotBound Set VName
_ Size
d = Size
d
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')